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..5595648 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]) ] @@ -163,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: @@ -172,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. diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py index bcc6357..bb0122e 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, @@ -11,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 @@ -224,6 +226,82 @@ 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 + 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."""