disable __setitem__ in static mode & add API paddle.static.setitem wi…

…th dy2st strategy (PaddlePaddle#53682)

* add paddle.static.setitem

* add some help doc

* support setitem

* support machanism

* add more unittest

* remove usless code

* raise error in static setitem

* fix d2s UT

* remove static only for both-used code

* fix bool set_value in static, fix set_value_op UT

* fix unittests

* [May case some error]: remove inplace-version check

* add two test case for dy2st

* fix function in vision

* fix dy2st setitem support, refine UT case

* fix slice in static_mode

* add ParametersMap

* remove pop

* modify place

* [fix]: variable is also a tensor

* rewrite some ut & remove slicetransformer in dy2st

* solve error in static-mode

* fix ut

* return a result for set_array_write

* fix test_set_value_op_xpu

* code is different in dynamic / static mode

---------

Co-authored-by: Aurelius84 <[email protected]>
Co-authored-by: NotHaozi <[email protected]>

python/paddle/distribution/multinomial.py

@@ -130,7 +130,12 @@ def log_prob(self, value):
         logits, value = paddle.broadcast_tensors(
             [paddle.log(self.probs), value]
         )
-        logits[(value == 0) & (paddle.isinf(logits))] = 0
+        if paddle.in_dynamic_mode():
+            logits[(value == 0) & (paddle.isinf(logits))] = 0
+        else:
+            logits = paddle.static.setitem(
+                logits, (value == 0) & (paddle.isinf(logits)), 0
+            )
 
         return (
             paddle.lgamma(value.sum(-1) + 1)

python/paddle/fluid/framework.py

@@ -2295,7 +2295,14 @@ def __getitem__(self, item):
         return _getitem_impl_(self, item)
 
     def __setitem__(self, item, value):
-        return _setitem_impl_(self, item, value)
+        from .dygraph.base import in_declarative_mode
+
+        if in_declarative_mode():
+            return _setitem_impl_(self, item, value)
+        else:
+            raise RuntimeError(
+                "In static mode, the __setitem__ (looks like: x[indices] = values) should not be used. Please use x = paddle.static.setitem(x, indices, values)"
+            )
 
     def get_value(self, scope=None):
         """

python/paddle/fluid/variable_index.py

@@ -215,7 +215,9 @@ def set_item(self, tensor_origin, value):
         out = paddle.scatter(t_reshape, index_1d, value_1d)
         if tensor_type is not None:
             out = out.astype(tensor_type)
-        tensor_origin[:] = out.reshape(tensor_origin.shape)
+        tensor_origin = _setitem_impl_(
+            tensor_origin, ..., out.reshape(tensor_origin.shape)
+        )
 
         return tensor_origin
 
@@ -617,7 +619,7 @@ def _setitem_for_tensor_array(var, item, value):
     If item is case (1), we perform paddle.tensor.array_write,
     in other cases, we raise a NotImplementedError.
     """
-    from ..framework import LayerHelper, core
+
     from .framework import Variable
 
     assert (
@@ -632,7 +634,7 @@ def _setitem_for_tensor_array(var, item, value):
 
         item = paddle.cast(to_static_variable(item), dtype='int64')
         value = to_static_variable(value)
-        array_write(x=value, i=item, array=var)
+        return array_write(x=value, i=item, array=var)
     else:
         raise NotImplementedError(
             "Only support __setitem__ by Int/Variable in tensor_array, but gets {}".format(
@@ -807,17 +809,31 @@ def _setitem_impl_(var, item, value):
 
     if paddle.in_dynamic_mode():
         var._bump_inplace_version()
+        output = var
+    else:
+        helper = paddle.fluid.layer_helper.LayerHelper('set_value', **locals())
+        output = helper.create_variable_for_type_inference(dtype=var.dtype)
 
     cur_block = default_main_program().current_block()
     cur_block.append_op(
         type="set_value",
         inputs=inputs,
-        outputs={'Out': var},
+        outputs={'Out': output},
         attrs=attrs,
         inplace_map={"Input": "Out"},
     )
 
-    return var
+    if not paddle.in_dynamic_mode():
+        # map var to the new output
+        from paddle.jit.dy2static.program_translator import (
+            ProgramTranslator,
+        )
+
+        ProgramTranslator.get_instance()._params_map.add(
+            cur_block.program, var.desc.id(), output
+        )
+
+    return output
 
 
 # the item is a tensor of bool
@@ -848,11 +864,19 @@ def idx_not_empty(var, item, value):
         gather_val = gather_nd(var, idx)
         gather_val_new = value - gather_val
         out = scatter_nd_add(var, idx, gather_val_new)
-        var[:] = out
+        var = _setitem_impl_(var, ..., out)
+        return var
+
+    def idx_is_empty(var):
+        return var
 
     from paddle.static.nn import cond
 
     # If all the bool index is False, just do nothing
-    cond(item.any(), lambda: idx_not_empty(var, item, value))
+    var = cond(
+        item.any(),
+        lambda: idx_not_empty(var, item, value),
+        lambda: idx_is_empty(var),
+    )
 
     return var

python/paddle/incubate/optimizer/functional/lbfgs.py

@@ -178,16 +178,23 @@ def body(
             shape=[], fill_value=(head - 1).mod(history_size), dtype='int64'
         )
 
-        def cond(i, q):
+        def cond(i, q, ai_vec):
             return i != tail
 
-        def body(i, q):
-            ai_vec[i] = rhok_vec[i] * paddle.dot(sk_vec[i], q)
+        def body(i, q, ai_vec):
+            if paddle.in_dynamic_mode():
+                ai_vec[i] = rhok_vec[i] * paddle.dot(sk_vec[i], q)
+            else:
+                ai_vec = paddle.static.setitem(
+                    ai_vec, i, rhok_vec[i] * paddle.dot(sk_vec[i], q)
+                )
             q = q - ai_vec[i] * yk_vec[i]
             i = (i - 1).mod(history_size)
-            return i, q
+            return i, q, ai_vec
 
-        paddle.static.nn.while_loop(cond=cond, body=body, loop_vars=[i, q])
+        paddle.static.nn.while_loop(
+            cond=cond, body=body, loop_vars=[i, q, ai_vec]
+        )
 
         r = paddle.matmul(H0, q)
 
@@ -234,10 +241,14 @@ def body(i, r):
             lambda: paddle.full(shape=[1], fill_value=1000.0, dtype=dtype),
             lambda: 1.0 / rhok_inv,
         )
-
-        sk_vec[head] = sk
-        yk_vec[head] = yk
-        rhok_vec[head] = rhok
+        if paddle.in_dynamic_mode():
+            sk_vec[head] = sk
+            yk_vec[head] = yk
+            rhok_vec[head] = rhok
+        else:
+            sk_vec = paddle.static.setitem(sk_vec, head, sk)
+            yk_vec = paddle.static.setitem(yk_vec, head, yk)
+            rhok_vec = paddle.static.setitem(rhok_vec, head, rhok)
         head = (head + 1) % history_size
 
         def true_fn(tail):

python/paddle/jit/dy2static/convert_operators.py

@@ -20,7 +20,7 @@
     _convert_into_variable,
     in_declarative_mode,
 )
-from paddle.fluid.framework import Variable, core
+from paddle.fluid.framework import Variable, core, default_main_program
 from paddle.fluid.layers import control_flow
 from paddle.fluid.layers.control_flow import while_loop
 
@@ -48,6 +48,19 @@ def convert_load(x):
         TODO:(@xiongkun) may run convert_load in dygraph mode, which should be fixed.
         """
         return _convert_into_variable(x)
+
+    # get the new output of the var
+    if in_declarative_mode() and isinstance(x, Variable):
+        cur_block = default_main_program().current_block()
+
+        from paddle.jit.dy2static.program_translator import ProgramTranslator
+
+        new_var = ProgramTranslator.get_instance()._params_map.get(
+            cur_block.program, x.desc.id()
+        )
+        if new_var is not None:
+            return new_var
+
     return x
 
 

python/paddle/jit/dy2static/program_translator.py

@@ -1125,6 +1125,36 @@ def _program_hash(self, program):
         return id(program)
 
 
+class ParametersMap:
+    def __init__(self):
+        self.params_dict = {}
+
+    @synchronized
+    def add(self, program, id, param):
+        """use the default_program as key, append param the parameter list."""
+        key = self._program_hash(program)
+        if key not in self.params_dict:
+            self.params_dict[key] = {}
+
+        params = self.params_dict[key]
+        params[id] = param
+
+    def get(self, program, id):
+        params = self.params_dict.get(self._program_hash(program))
+        if params is None:
+            return None
+        if id in params.keys():
+            return params[id]
+        return None
+
+    def _program_hash(self, program):
+        """
+        because program is not deleted while calling from_func_spec.
+        so it's ok to use id(program)
+        """
+        return id(program)
+
+
 class FallbackProgramLayer:
     __slots__ = [
         '_instance',
@@ -1386,6 +1416,7 @@ def __init__(self):
         self._initialized = True
         self._program_cache = ProgramCache()
         self._params_recorder = ParametersRecorder()
+        self._params_map = ParametersMap()
         self.enable_to_static = True
 
     def enable(self, enable_to_static):

python/paddle/static/__init__.py

@@ -37,6 +37,7 @@
 from ..fluid import Scope  # noqa: F401
 from .input import data  # noqa: F401
 from .input import InputSpec  # noqa: F401
+from .input import setitem  # noqa: F401
 
 from ..tensor.creation import create_parameter  # noqa: F401
 from ..tensor.creation import create_global_var  # noqa: F401

python/paddle/static/input.py

@@ -18,6 +18,8 @@
 from paddle.fluid.framework import convert_np_dtype_to_dtype_, static_only
 from paddle.fluid.layer_helper import LayerHelper
 
+from ..fluid.variable_index import _setitem_impl_
+
 __all__ = []
 
 
@@ -342,3 +344,28 @@ def __eq__(self, other):
 
     def __ne__(self, other):
         return not self == other
+
+
+def setitem(x, index, value):
+    """
+    x(Tensor): input Tensor.
+    index(Scalar|Tuple|List|Tensor): Where should be set value.
+    value(Scalar|Tensor): The value which is going to be set.
+
+    [How to write index?]
+    1. ':' -> slice(),
+       (1) a[:]=v -> setitem(a, slice(None,None,None), v)
+       (2) a[1::2] -> setitem(a, slice(1,None,2), v)
+
+    2. if there are multiple indexes for axes, use TUPLE (Not LIST) to pack them.
+       (1) a[1, 2]=v -> setitem(a, (1, 2), v)
+       (2) a[[1,2],[2,3]]=v -> setitem(a, ([1,2],[2,3]), v)
+       (3) a[1,:, 3] = v -> setitem(a, (1, slice(None,None,None),3), v)
+       (4) a[1, ..., 2]=v -> setitem(a, (1, ..., 2), v)
+
+    3. You can always use TUPLE as index input， even there is only one index.
+       (1) a[Tensor([10,10])]=v -> setitem(a, (Tensor([10,10]),), v)
+       (2) a[1] = v -> setitem(a, (1,), v)
+    """
+
+    return _setitem_impl_(x, index, value)

python/paddle/tensor/math.py

@@ -5788,11 +5788,26 @@ def vander(x, n=None, increasing=False, name=None):
 
     res = paddle.empty([x.shape[0], n], dtype=x.dtype)
 
-    if n > 0:
-        res[:, 0] = paddle.to_tensor([1], dtype=x.dtype)
-    if n > 1:
-        res[:, 1:] = x[:, None]
-        res[:, 1:] = paddle.cumprod(res[:, 1:], dim=-1)
+    if paddle.in_dynamic_mode():
+        if n > 0:
+            res[:, 0] = paddle.to_tensor([1], dtype=x.dtype)
+        if n > 1:
+            res[:, 1:] = x[:, None]
+            res[:, 1:] = paddle.cumprod(res[:, 1:], dim=-1)
+    else:
+        if n > 0:
+            res = paddle.static.setitem(
+                res, (slice(None), 0), paddle.to_tensor([1], dtype=x.dtype)
+            )
+        if n > 1:
+            res = paddle.static.setitem(
+                res, (slice(None), slice(1, None)), x[:, None]
+            )
+            res = paddle.static.setitem(
+                res,
+                (slice(None), slice(1, None)),
+                paddle.cumprod(res[:, 1:], dim=-1),
+            )
     res = res[:, ::-1] if not increasing else res
     return res