Allow inplace ScalarLoop

Author: ricardoV94

pytensor/scalar/loop.py

@@ -218,11 +218,11 @@ def c_code_template(self):
             c: f"%(i{int(i)})s"
             for i, c in enumerate(fgraph.inputs[n_update:], start=n_update + 1)
         }
-        update_subd = {
+        out_subd = {
             u: f"%(o{int(i)})s" for i, u in enumerate(fgraph.outputs[:n_update])
         }
         until_subd = {u: "until" for u in fgraph.outputs[n_update:]}
-        subd = {**carry_subd, **constant_subd, **update_subd, **until_subd}
+        subd = {**carry_subd, **constant_subd, **until_subd}
 
         for var in fgraph.variables:
             if var.owner is None:
@@ -246,11 +246,11 @@ def c_code_template(self):
             _c_code += "bool until = 1;\n\n"
 
         # Copy carried inputs
-        for i, (var, name) in enumerate(carry_subd.items()):
-            copy_var_name = f"{name}_copy{i}"
-            _c_code += f"{var.type.dtype_specs()[1]} {copy_var_name} = {name};\n"
-            carry_subd[var] = copy_var_name
-            subd[var] = copy_var_name
+        for i, (var, name) in enumerate(carry_subd.items(), start=1):
+            carry_var_name = f"{name}_carry{i}"
+            _c_code += f"{var.type.dtype_specs()[1]} {carry_var_name} = {name};\n"
+            carry_subd[var] = carry_var_name
+            subd[var] = carry_var_name
 
         # _c_code += 'printf("inputs=[");'
         # for i in range(1, len(fgraph.inputs)):
@@ -259,9 +259,8 @@ def c_code_template(self):
 
         _c_code += "\nfor(%(n_steps_dtype)s i = 0; i < %(n_steps)s; i++){\n"
 
-        self.nodenames = [
-            f"%(nodename)s_subnode{int(j)}" for j, n in enumerate(fgraph.toposort())
-        ]
+        # Used by self.c_support_code_apply
+        self.nodenames = nodenames = []
 
         i = 0
         for j, node in enumerate(fgraph.toposort()):
@@ -271,9 +270,13 @@ def c_code_template(self):
                     name = f"V%(id)s_tmp{int(i)}"
                     subd[output] = name
                     _c_code += f"{output.type.dtype_specs()[1]} {name};\n"
+
+            nodename = f"%(nodename)s_subnode{int(j)}"
+            nodenames.append(nodename)
+
             s = node.op.c_code(
                 node,
-                self.nodenames[j],
+                nodename,
                 # Any node that depended on `init` will depend on `update` instead
                 # The initial value of `update` was set to `init` before the loop
                 [subd[input] for input in node.inputs],
@@ -283,10 +286,12 @@ def c_code_template(self):
             _c_code += s
             _c_code += "\n"
 
-        # Set the carry variables to the output variables
+        # Update the carry variables to the output variables
         _c_code += "\n"
-        for init, update in zip(carry_subd.values(), update_subd.values(), strict=True):
-            _c_code += f"{init} = {update};\n"
+        for carry, out in zip(
+            carry_subd.values(), fgraph.outputs[:n_update], strict=True
+        ):
+            _c_code += f"{carry} = {subd[out]};\n"
 
         # _c_code += 'printf("%%ld\\n", i);\n'
         # for carry in range(1, 10):
@@ -298,6 +303,10 @@ def c_code_template(self):
         # End of the loop
         _c_code += "}\n"
 
+        # Assign the carry variables to the outputs
+        for out, carry in zip(out_subd.values(), carry_subd.values(), strict=True):
+            _c_code += f"{out} = {carry};\n"
+
         # Output until flag
         if self.is_while:
             _c_code += f"%(o{len(fgraph.outputs)-1})s = until;\n"
@@ -332,4 +341,4 @@ def c_code(self, node, nodename, inames, onames, sub):
         return res
 
     def c_code_cache_version_outer(self):
-        return (3,)
+        return (4,)

tests/scalar/test_loop.py

@@ -19,7 +19,7 @@
 )
 from pytensor.scalar.loop import ScalarLoop
 from pytensor.tensor import exp as tensor_exp
-from pytensor.tensor import vector
+from pytensor.tensor import lvector
 from pytensor.tensor.elemwise import Elemwise
 
 
@@ -260,35 +260,43 @@ def test_inner_loop(mode):
     )
 
 
-def test_elemwise_inplace():
-    x0 = float64("x0")
-    y0 = float64("y0")
-    x = x0 - y0
-    y = y0 - x0
-    op = Elemwise(ScalarLoop(init=[x0, y0], constant=[], update=[x, y]))
-
-    n_steps = vector("n_steps", dtype="int64")
-    x0v = vector("x0")
-    y0v = vector("y0")
-    xv, yv = op(n_steps, x0v, y0v)
+@pytest.mark.parametrize("mutate_arg_idx", (0, 1, 2, 3))
+def test_elemwise_inplace(mutate_arg_idx):
+    x0 = int64("x0")
+    y0 = int64("y0")
+    c = int64("c")
+    x = x0 - y0 + c
+    y = y0 - x0 + c
+    op = Elemwise(ScalarLoop(init=[x0, y0], constant=[c], update=[x, y]))
+
+    n_steps = lvector("n_steps")
+    x0v = lvector("x0")
+    y0v = lvector("y0")
+    cv = lvector("c")
+    xv, yv = op(n_steps, x0v, y0v, cv)
+
+    inputs = [
+        In(inp, mutable=i == mutate_arg_idx)
+        for i, inp in enumerate([n_steps, x0v, y0v, cv])
+    ]
 
     fn = function(
-        [In(n_steps, mutable=True), In(x0v, mutable=True), In(y0v, mutable=True)],
+        inputs,
         [xv, yv],
         mode=get_default_mode().including("inplace"),
     )
     elem_op = fn.maker.fgraph.outputs[0].owner.op
     assert isinstance(elem_op, Elemwise) and isinstance(elem_op.scalar_op, ScalarLoop)
     destroy_map = elem_op.destroy_map
-    assert destroy_map in ({0: [1], 1: [2]}, {0: [2], 1: [2]})
+    assert destroy_map == {0: [mutate_arg_idx]}
 
-    n_test = np.array([1, 4, 8], dtype="int32")
-    x0v_test = np.array([0, 0, 0], dtype=x0v.dtype)
-    y0v_test = np.array([1, 1, 1], dtype=y0v.dtype)
+    n_test = np.array([1, 4, 8], dtype="int64")
+    x0v_test = np.array([0, 0, 0], dtype="int64")
+    y0v_test = np.array([1, 1, 1], dtype="int64")
+    cv_test = np.array([0, 0, 0], dtype="int64")
 
-    xv_res, yv_res = fn(n_test, x0v_test, y0v_test)
+    xv_res, yv_res = fn(n_test, x0v_test, y0v_test, cv_test)
     # Check the outputs are the destroyed inputs
-    assert xv_res is (x0v_test, y0v_test)[destroy_map[0][0] - 1]
-    assert yv_res is (x0v_test, y0v_test)[destroy_map[1][0] - 1]
+    assert xv_res is (n_test, x0v_test, y0v_test, cv_test)[mutate_arg_idx]
     np.testing.assert_allclose(xv_res, [-1, -8, -128])
     np.testing.assert_allclose(yv_res, [1, 8, 128])