From d4f2307cfb5b982e625c4a9c6c81fc1482c16f4d Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 13 Aug 2026 17:07:11 +0200 Subject: [PATCH 1/2] Fold an Alloc whose inplace client is itself constant --- pytensor/tensor/basic.py | 5 +++++ tests/tensor/rewriting/test_basic.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 52f63d9e34..de1267d6e2 100644 --- a/pytensor/tensor/basic.py +++ b/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 diff --git a/tests/tensor/rewriting/test_basic.py b/tests/tensor/rewriting/test_basic.py index 7943507472..6a25ff867a 100644 --- a/tests/tensor/rewriting/test_basic.py +++ b/tests/tensor/rewriting/test_basic.py @@ -732,6 +732,25 @@ 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) + @pytest.mark.xfail( reason="PyTensor rewrites constants before stabilization. " "This breaks stabilization rewrites in some cases. See #504.", From 2899a9503765b332ed813cdace47cefc2cb168ac Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 14 Aug 2026 11:47:36 +0200 Subject: [PATCH 2/2] Fold an AllocEmpty whose clients are themselves constants --- pytensor/tensor/basic.py | 11 ++++++++++- tests/tensor/rewriting/test_basic.py | 20 ++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index de1267d6e2..b846998194 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -4284,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] diff --git a/tests/tensor/rewriting/test_basic.py b/tests/tensor/rewriting/test_basic.py index 6a25ff867a..2fc7da48cc 100644 --- a/tests/tensor/rewriting/test_basic.py +++ b/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, @@ -751,6 +752,25 @@ def test_alloc_folds_when_its_inplace_client_is_constant(self, advanced): 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.",