From bb0269f6b0607ca93a383921ac1d278de2904db2 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 11 Aug 2026 21:08:29 -0500 Subject: [PATCH 1/3] Stop the frozen-generator check from firing on inner-graph inputs An op carrying an inner graph may draw from a generator it receives or merely accept it, and only that inner graph knows which, so treating every non-output use as a draw rejected graphs that never draw at all. `collect_default_updates` already raises for an inner graph that draws without threading the generator back out. --- pytensor_ml/pytensorf/compile.py | 4 ++-- pytensor_ml/pytensorf/rng.py | 11 ++++++----- tests/test_pytensorf.py | 13 +++++++++++++ 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/pytensor_ml/pytensorf/compile.py b/pytensor_ml/pytensorf/compile.py index 8be9c26..cccb310 100644 --- a/pytensor_ml/pytensorf/compile.py +++ b/pytensor_ml/pytensorf/compile.py @@ -81,8 +81,8 @@ def function( if frozen: raise ValueError( f"The graph draws from {[str(generator.name or generator) for generator in frozen]}, which " - "nothing advances, so every call would repeat the same values. A generator read by two " - "different draws has no single next state, which is the usual cause. Give each draw its own " + "nothing advances, so every call would repeat the same values. Two draws off one generator is " + "the cause: it has no single next state, so none can be derived. Give each draw its own " "generator, thread one through with `next_rng, draw = pt.random.normal(rng=rng, " "return_next_rng=True)`, or pass an update for it yourself." ) diff --git a/pytensor_ml/pytensorf/rng.py b/pytensor_ml/pytensorf/rng.py index 54e29a4..0a19bc6 100644 --- a/pytensor_ml/pytensorf/rng.py +++ b/pytensor_ml/pytensorf/rng.py @@ -50,10 +50,11 @@ def find_generators_drawn_from( outputs: Sequence[Variable], ) -> list[RandomGeneratorSharedVariable]: """ - Return the shared generators the graph draws from, excluding any it only hands back. + Return the shared generators a draw op in this graph consumes. - A generator returned as an output is read without being consumed, so it needs no update; one an op draws - from does, and a graph that draws from a generator nothing advances repeats the same values forever. + Being read is not being consumed: a generator handed back as an output needs no update, and neither does + one passed into an op carrying an inner graph, which may draw from it or merely accept it. Only that + inner graph knows, and reading inner graphs is :func:`collect_default_updates`' job. Parameters ---------- @@ -63,14 +64,14 @@ def find_generators_drawn_from( Returns ------- list of RandomGeneratorSharedVariable - The generators some node in the graph draws from, in graph-input order. + The generators a draw op consumes, in graph-input order. """ fgraph = FunctionGraph(outputs=list(outputs), clone=False) return [ generator for generator in fgraph.inputs if isinstance(generator, RandomGeneratorSharedVariable) - and any(not isinstance(client.op, Output) for client, _ in fgraph.clients[generator]) + and any(isinstance(client.op, RNGConsumerOp) for client, _ in fgraph.clients[generator]) ] diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py index bcc6357..bbb815b 100644 --- a/tests/test_pytensorf.py +++ b/tests/test_pytensorf.py @@ -3,6 +3,7 @@ import pytest from pytensor import config, shared +from pytensor.compile.builders import OpFromGraph from pytensor.gradient import ( DisconnectedInputError, disconnected_grad, @@ -224,6 +225,18 @@ def test_a_shared_generator_is_accepted_when_the_caller_advances_it(): assert float(first[1]) != float(second[1]) # and so does the one off the threaded next state +def test_a_generator_an_inner_graph_never_draws_from_is_left_alone(): + """Receiving a generator is not drawing from one, and only the inner graph knows which happened. Treating + every non-output use as a draw rejected this graph with a message saying it draws from a generator that + nothing in it touches.""" + rng = shared(np.random.default_rng(0), name="rng") + inner_rng = rng.type() + inner_x = pt.vector("x") + doubled = OpFromGraph([inner_rng, inner_x], [inner_x * 2.0])(rng, pt.ones(3)) + + np.testing.assert_allclose(function([], doubled)(), np.full(3, 2.0)) + + def test_a_draw_shared_between_an_output_and_an_update_is_rejected(): """An update expression is a reader like any other, so a generator drawn from by both an output and an update is read twice and cannot be advanced once.""" From 1c5a9cf01dbbee188f9df2a9432acfd35c6d81ec Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 11 Aug 2026 21:09:22 -0500 Subject: [PATCH 2/3] Test RNG handling for draws inside scan and OpFromGraph --- tests/test_pytensorf.py | 65 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py index bbb815b..bb0122e 100644 --- a/tests/test_pytensorf.py +++ b/tests/test_pytensorf.py @@ -12,6 +12,7 @@ zero_grad, ) from pytensor.graph.traversal import ancestors +from pytensor.scan import scan from pytensor.tensor import random as ptr from pytensor_ml.layers import Dropout, DropoutLayer, Linear, Sequential @@ -225,6 +226,70 @@ def test_a_shared_generator_is_accepted_when_the_caller_advances_it(): assert float(first[1]) != float(second[1]) # and so does the one off the threaded next state +def test_a_draw_inside_a_scan_advances_its_generator(): + """Nothing is passed to `function` as an update: a scan carrying its generator as a recurrent state + exposes it as an outer input with a matching outer output, and the collector reads that mapping. This is + the whole of the inner-graph handling, inherited from pymc, and the first test to run it.""" + rng = shared(np.random.default_rng(0), name="rng") + + def one_step(total, generator): + next_generator, draw = ptr.normal(size=(), rng=generator, return_next_rng=True) + return total + draw, next_generator + + # Two recurrent states out, not a graph and an updates dict -- the generator is carried as state. + trace, _generators = scan( + one_step, outputs_info=[pt.zeros(()), rng], n_steps=5, return_updates=False + ) + step = function([], trace) + + draws_within_one_call = np.diff(step(), prepend=0.0) + + # Both halves matter: a loop reusing one generator state draws the same number five times while its + # outer state still advances between calls, so neither assertion catches that alone. + assert len(np.unique(draws_within_one_call)) == 5 + assert not np.array_equal(step(), step()) + + +def test_a_scan_that_draws_without_threading_its_generator_is_rejected(): + """The loop draws from a generator it captured rather than carried, so there is no outer output holding + its final state and no way to advance it -- every call would replay the same five draws.""" + rng = shared(np.random.default_rng(0), name="rng") + + def one_step(total): + _, draw = ptr.normal(size=(), rng=rng, return_next_rng=True) + return total + draw + + trace = scan(one_step, outputs_info=[pt.zeros(())], n_steps=5, return_updates=False) + + with pytest.raises(ValueError, match="No update found for at least one RNG used in Scan"): + function([], trace[-1]) + + +def test_a_draw_inside_an_op_from_graph_advances_its_generator(): + """Every layer in this library is an OpFromGraph, so a layer that draws inside its own inner graph lands + here. Handing the advanced generator back as an output is what makes it reachable from outside.""" + rng = shared(np.random.default_rng(0), name="rng") + inner_rng = rng.type() + next_rng, inner_draw = ptr.normal(size=(3,), rng=inner_rng, return_next_rng=True) + draw, _ = OpFromGraph([inner_rng], [inner_draw, next_rng])(rng) + + step = function([], draw) + + assert not np.array_equal(step(), step()) + + +def test_an_op_from_graph_that_draws_without_threading_is_rejected(): + rng = shared(np.random.default_rng(0), name="rng") + inner_rng = rng.type() + _, inner_draw = ptr.normal(size=(3,), rng=inner_rng, return_next_rng=True) + draw = OpFromGraph([inner_rng], [inner_draw])(rng) + + with pytest.raises( + ValueError, match="No update found for at least one RNG used in OpFromGraph" + ): + function([], draw) + + def test_a_generator_an_inner_graph_never_draws_from_is_left_alone(): """Receiving a generator is not drawing from one, and only the inner graph knows which happened. Treating every non-output use as a draw rejected this graph with a message saying it draws from a generator that From e78adbd359c38dae53acdbe399c22e2411ba6d0f Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 11 Aug 2026 21:09:36 -0500 Subject: [PATCH 3/3] Say how to fix the inner-graph RNG errors --- pytensor_ml/pytensorf/rng.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pytensor_ml/pytensorf/rng.py b/pytensor_ml/pytensorf/rng.py index 0a19bc6..5595648 100644 --- a/pytensor_ml/pytensorf/rng.py +++ b/pytensor_ml/pytensorf/rng.py @@ -164,7 +164,9 @@ def find_default_update(clients, rng: Variable) -> None | Variable: next_rng = client.outputs[output_index] else: raise ValueError( - f"No update found for at least one RNG used in Scan Op {client_op}." + f"No update found for at least one RNG used in Scan Op {client_op}. Call " + "`collect_default_updates` inside the scan function and return what it gives you " + "as that step's updates." ) case OpFromGraph(): try: @@ -173,7 +175,9 @@ def find_default_update(clients, rng: Variable) -> None | Variable: return None except ValueError as exc: raise ValueError( - f"No update found for at least one RNG used in OpFromGraph Op {client_op}." + f"No update found for at least one RNG used in OpFromGraph Op {client_op}. Add " + "the advanced generator to the op's outputs, which " + "`pt.random.normal(rng=rng, return_next_rng=True)` gives you alongside the draw." ) from exc case _: # Unknown consumer; the caller must provide an update manually.