Allow inplace of Elemwise ScalarLoop

Author: ricardoV94

pytensor/scalar/basic.py

@@ -4102,6 +4102,7 @@ class ScalarInnerGraphOp(ScalarOp, HasInnerGraph):
 
     def __init__(self, *args, **kwargs):
         self.prepare_node_called = set()
+        super().__init__(*args, **kwargs)
 
     def _cleanup_graph(self, inputs, outputs):
         # TODO: We could convert to TensorVariable, optimize graph,

pytensor/scalar/loop.py

@@ -55,6 +55,7 @@ def __init__(
         constant: Sequence[Variable] | None = None,
         until: Variable | None = None,
         name="ScalarLoop",
+        **kwargs,
     ):
         if constant is None:
             constant = []
@@ -75,7 +76,7 @@ def __init__(
         self.nout = len(self.outputs)
         self.name = name
 
-        super().__init__()
+        super().__init__(**kwargs)
 
     def output_types(self, input_types):
         return self.outputs_type
@@ -115,7 +116,7 @@ def fgraph(self):
         self._fgraph = fgraph
         return self._fgraph
 
-    def clone(self):
+    def clone(self, name=None, **kwargs):
         if self.is_while:
             *update, until = self.outputs
         else:
@@ -127,28 +128,16 @@ def clone(self):
             update=update,
             constant=constant,
             until=until,
-            name=self.name,
+            name=self.name if name is None else name,
+            **kwargs,
         )
 
     @property
     def fn(self):
         raise NotImplementedError
 
     def make_new_inplace(self, output_types_preference=None, name=None):
-        """
-        This op.__init__ fct don't have the same parameter as other scalar op.
-        This break the insert_inplace_optimizer optimization.
-        This fct allow fix patch this.
-
-        """
-        d = {k: getattr(self, k) for k in self.init_param}
-        out = self.__class__(**d)
-        if name:
-            out.name = name
-        else:
-            name = out.name
-        super(ScalarLoop, out).__init__(output_types_preference, name)
-        return out
+        return self.clone(output_types_preference=output_types_preference, name=name)
 
     def make_node(self, n_steps, *inputs):
         assert len(inputs) == self.nin - 1

pytensor/tensor/rewriting/elemwise.py

@@ -24,7 +24,6 @@
 )
 from pytensor.graph.rewriting.db import SequenceDB
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined
-from pytensor.scalar.loop import ScalarLoop
 from pytensor.scalar.math import Grad2F1Loop, _grad_2f1_loop
 from pytensor.tensor.basic import (
     MakeVector,
@@ -74,15 +73,6 @@ def print_profile(cls, stream, prof, level=0):
             for n in sorted(ndim):
                 print(blanc, n, ndim[n], file=stream)
 
-    def candidate_input_idxs(self, node):
-        # TODO: Implement specialized InplaceCompositeOptimizer with logic
-        #  needed to correctly assign inplace for multi-output Composites
-        #  and ScalarLoops
-        if isinstance(node.op.scalar_op, ScalarLoop):
-            return []
-        else:
-            return range(len(node.outputs))
-
     def apply(self, fgraph):
         r"""
 
@@ -173,7 +163,7 @@ def apply(self, fgraph):
 
                 baseline = op.inplace_pattern
                 candidate_outputs = [
-                    i for i in self.candidate_input_idxs(node) if i not in baseline
+                    i for i in range(len(node.outputs)) if i not in baseline
                 ]
                 # node inputs that are Constant, already destroyed,
                 # or fgraph protected inputs and fgraph outputs can't be used as
@@ -190,7 +180,7 @@ def apply(self, fgraph):
                 ]
             else:
                 baseline = []
-                candidate_outputs = self.candidate_input_idxs(node)
+                candidate_outputs = range(len(node.outputs))
                 # node inputs that are Constant, already destroyed,
                 # fgraph protected inputs and fgraph outputs can't be used as inplace
                 # target.

tests/scalar/test_loop.py

@@ -3,7 +3,8 @@
 import numpy as np
 import pytest
 
-from pytensor import Mode, function
+from pytensor import In, Mode, function
+from pytensor.compile import get_default_mode
 from pytensor.scalar import (
     Composite,
     as_scalar,
@@ -18,6 +19,8 @@
 )
 from pytensor.scalar.loop import ScalarLoop
 from pytensor.tensor import exp as tensor_exp
+from pytensor.tensor import vector
+from pytensor.tensor.elemwise import Elemwise
 
 
 mode = pytest.mark.parametrize(
@@ -255,3 +258,37 @@ def test_inner_loop(mode):
         out16,
         3**2 + 2.5,
     )
+
+
+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)
+
+    fn = function(
+        [In(n_steps, mutable=True), In(x0v, mutable=True), In(y0v, mutable=True)],
+        [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]})
+
+    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)
+
+    xv_res, yv_res = fn(n_test, x0v_test, y0v_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]
+    np.testing.assert_allclose(xv_res, [-1, -8, -128])
+    np.testing.assert_allclose(yv_res, [1, 8, 128])