Do not skip validation between consecutive Elemwise inplace replacements

Author: ricardoV94

pytensor/graph/destroyhandler.py

@@ -7,7 +7,6 @@
 import itertools
 from collections import deque
 
-import pytensor
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Constant
 from pytensor.graph.features import AlreadyThere, Bookkeeper
@@ -223,7 +222,7 @@ def _build_droot_impact(destroy_handler):
     return droot, impact, root_destroyer
 
 
-def fast_inplace_check(fgraph, inputs):
+def inplace_candidates(fgraph, inputs, protected_inputs=None):
     """
     Return the variables in inputs that are possible candidate for as inputs of
     inplace operation.
@@ -234,22 +233,40 @@ def fast_inplace_check(fgraph, inputs):
         Inputs Variable that you want to use as inplace destination.
 
     """
-    Supervisor = pytensor.compile.function.types.Supervisor
-    protected_inputs = list(
-        itertools.chain.from_iterable(
-            f.protected for f in fgraph._features if isinstance(f, Supervisor)
+
+    def view_is_protected(view_dict, variable, protected):
+        if variable in protected:
+            return True
+        try:
+            while True:
+                variable = view_dict[variable]
+                if variable in protected:
+                    return True
+        except KeyError:
+            return False
+
+    if protected_inputs is None:
+        from pytensor.compile.function.types import Supervisor
+
+        protected_inputs = set(
+            itertools.chain.from_iterable(
+                f.protected for f in fgraph._features if isinstance(f, Supervisor)
+            )
+        )
+        protected_inputs.update(fgraph.outputs)
+
+    has_destroyers = fgraph.has_destroyers
+    view_i = fgraph.destroy_handler.view_i
+
+    return [
+        inp
+        for inp in inputs
+        if (
+            not isinstance(inp, Constant)
+            and not view_is_protected(view_i, inp, protected_inputs)
+            and not has_destroyers([inp])
         )
-    )
-    protected_inputs.extend(fgraph.outputs)
-
-    inputs = [
-        i
-        for i in inputs
-        if not isinstance(i, Constant)
-        and not fgraph.has_destroyers([i])
-        and i not in protected_inputs
     ]
-    return inputs
 
 
 class DestroyHandler(Bookkeeper):

pytensor/tensor/rewriting/blockwise.py

@@ -1,8 +1,6 @@
-import itertools
-
-from pytensor.compile import Supervisor
 from pytensor.compile.mode import optdb
 from pytensor.graph import Constant, node_rewriter
+from pytensor.graph.destroyhandler import inplace_candidates
 from pytensor.graph.replace import vectorize_node
 from pytensor.graph.rewriting.basic import copy_stack_trace, in2out, out2in
 from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
@@ -270,28 +268,20 @@ def blockwise_inplace(fgraph, node):
         # Op already has inplace
         return
 
-    # Find out valid inputs for inplacing
+    # Find out valid inputs for in-placing
     batch_ndim = blockwise_op.batch_ndim(node)
     out_batch_bcast = node.outputs[0].type.broadcastable[:batch_ndim]
+    inputs = node.inputs
 
-    protected_inputs = [
-        f.protected for f in fgraph._features if isinstance(f, Supervisor)
-    ]
-    protected_inputs = list(itertools.chain.from_iterable(protected_inputs))
-    protected_inputs.extend(fgraph.outputs)
     allowed_inplace_inputs = [
-        idx
-        for idx, inp in enumerate(node.inputs)
-        if
-        (
-            # Constants would need to be recreated every time if inplaced
-            not isinstance(inp, Constant)
-            # We can only inplace on inputs that are not being broadcasted
-            # As those are reused across iterations of Blockwise
-            and node.inputs[idx].type.broadcastable[:batch_ndim] == out_batch_bcast
-            # Inputs that are marked as protected or destroyed can't be inplaced
-            and not fgraph.has_destroyers([inp])
-            and inp not in protected_inputs
+        inputs.index(inp)
+        for inp in inplace_candidates(
+            fgraph,
+            [
+                inp
+                for inp in inputs
+                if inp.type.broadcastable[:batch_ndim] == out_batch_bcast
+            ],
         )
     ]