Fold an Alloc whose inplace client is itself constant

pytensor/tensor/basic.py

@@ -1798,6 +1798,11 @@ def do_constant_folding(self, fgraph, node):
                     client_op,
                     IncSubtensor | AdvancedIncSubtensor | Gemv | CGemv | Ger | CGer,
                 )
+                # ... unless the client is constant in every other input, in which
+                # case it folds away as well and no copy survives to run time.
+                # Refusing to fold then leaves the whole cone recomputing the same
+                # bytes on every call.
+                and not all(isinstance(inp, Constant) for inp in client.inputs[1:])
             ):
                 # Ops that will work inplace on the Alloc. So if they
                 # get constant_folded, they would copy the constant
@@ -4279,7 +4284,16 @@ def c_code_cache_version(self):
         return (4,)
 
     def do_constant_folding(self, fgraph, node):
-        return False
+        # The contents are undefined, so a constant is only ever as good as the
+        # buffer it replaces if it does not survive the fold: every client must be
+        # constant in its other inputs, so that it folds away as well.
+        [out] = node.outputs
+        clients = fgraph.clients[out]
+        return bool(clients) and all(
+            not isinstance(client.op, Output)
+            and all(inp is out or isinstance(inp, Constant) for inp in client.inputs)
+            for client, _ in clients
+        )
 
     def connection_pattern(self, node):
         return [[False] for i in node.inputs]

tests/tensor/rewriting/test_basic.py

@@ -25,6 +25,7 @@
 from pytensor.scalar import Composite, float64
 from pytensor.tensor.basic import (
     Alloc,
+    AllocEmpty,
     Join,
     MakeVector,
     ScalarFromTensor,
@@ -732,6 +733,44 @@ def test_constant_folding(self):
         assert len(topo) == 2
         assert all(isinstance(n.op, DeepCopyOp) for n in topo)
 
+    @pytest.mark.parametrize("advanced", [False, True], ids=["basic", "advanced"])
+    def test_alloc_folds_when_its_inplace_client_is_constant(self, advanced):
+        idx = constant(np.array([0, 2])) if advanced else slice(0, 2)
+
+        # An Alloc feeding an (Advanced)IncSubtensor is not folded when the update
+        # is not constant: the IncSubtensor would have to copy the constant.
+        y = pt.matrix("y")
+        out = pt.set_subtensor(pt.zeros((3, 4))[idx], y)
+        fg = FunctionGraph([y], [out], clone=True)
+        topo_constant_folding.apply(fg)
+        assert any(isinstance(node.op, Alloc) for node in fg.apply_nodes)
+
+        # It is folded when the IncSubtensor is itself constant: that folds away
+        # too, so no copy survives to run time.
+        out = pt.set_subtensor(pt.zeros((3, 4))[idx], np.ones((2, 4)))
+        fg = FunctionGraph([], [out], clone=True)
+        topo_constant_folding.apply(fg)
+        assert isinstance(fg.outputs[0], Constant)
+
+    def test_alloc_empty_folds_only_into_a_constant_client(self):
+        # An undefined buffer is worth keeping unless it folds away entirely.
+        y = pt.matrix("y")
+        out = pt.set_subtensor(pt.empty((3, 4))[:2], y)
+        fg = FunctionGraph([y], [out], clone=True)
+        topo_constant_folding.apply(fg)
+        assert any(isinstance(node.op, AllocEmpty) for node in fg.apply_nodes)
+
+        out = pt.set_subtensor(pt.empty((3, 4))[:2], np.ones((2, 4)))
+        fg = FunctionGraph([], [out], clone=True)
+        topo_constant_folding.apply(fg)
+        assert isinstance(fg.outputs[0], Constant)
+
+        # It is also not folded when it is an output, which would hand out the same
+        # undefined values on every call.
+        fg = FunctionGraph([], [pt.empty((3, 4))], clone=True)
+        topo_constant_folding.apply(fg)
+        assert any(isinstance(node.op, AllocEmpty) for node in fg.apply_nodes)
+
     @pytest.mark.xfail(
         reason="PyTensor rewrites constants before stabilization. "
         "This breaks stabilization rewrites in some cases. See #504.",