Stop the frozen-generator check from rejecting inner-graph inputs

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."
         )

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.

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."""