Fix RNG seeding and reject generators nothing advances

pytensor_ml/layers/dropout.py

@@ -6,6 +6,7 @@
 
 from pytensor import config
 from pytensor.compile.sharedvalue import shared
+from pytensor.tensor.random.variable import RandomGeneratorSharedVariable
 
 from pytensor_ml.base import Layer, UnaryLayerOp
 
@@ -18,17 +19,53 @@ def build_inner_graph(self, X, mask):
 
 
 class Dropout(Layer):
+    """
+    Zero a random fraction of its input, rescaling the rest so the expected sum is unchanged.
+
+    Parameters
+    ----------
+    name : str, optional
+        Name of the layer, used to name its output and its generators. Default "Dropout".
+    p : float
+        Probability of zeroing each element. Default 0.5.
+    random_state : int, Generator, or other seed, optional
+        Seed for the layer's masks. Draws are reproducible under a given seed, including when the layer is
+        applied at several points in a network. Seeded from fresh entropy when omitted.
+
+    Attributes
+    ----------
+    generators : list of RandomGeneratorSharedVariable
+        One generator per application of the layer, in the order they were applied. Set a value on these to
+        steer or restore the masks.
+    """
+
     def __init__(self, name: str | None = None, p: float = 0.5, random_state: Any | None = None):
         if p < 0.0 or p > 1.0:
             raise ValueError(f"Dropout probability has to be between 0 and 1, but got {p}")
         self.name = name if name else "Dropout"
         self.p = p
-        self.rng = shared(np.random.default_rng(random_state))
+        self.generators: list[RandomGeneratorSharedVariable] = []
+        self._generator_source = np.random.default_rng(random_state)
+
+    def _own_generator(self) -> RandomGeneratorSharedVariable:
+        """Return a generator for one application of this layer, spawned from its seed.
+
+        Each application draws its own mask, so each needs a generator of its own: a generator read by two
+        draws has no single next state, and nothing can advance it. Spawning keeps the whole layer
+        reproducible under ``random_state`` however many times it is applied.
+        """
+        generator = shared(
+            self._generator_source.spawn(1)[0], name=f"{self.name}/rng_{len(self.generators)}"
+        )
+        self.generators.append(generator)
+        return generator
 
     def __call__(self, X: pt.TensorLike) -> pt.TensorVariable:
         X = pt.as_tensor(X)
         p = pt.as_tensor(self.p, dtype=config.floatX)
-        _, mask = ptr.bernoulli(p=1 - p, size=X.shape, rng=self.rng, return_next_rng=True)
+        _, mask = ptr.bernoulli(
+            p=1 - p, size=X.shape, rng=self._own_generator(), return_next_rng=True
+        )
         mask = mask.astype(config.floatX)
 
         X_masked = DropoutLayer(

pytensor_ml/optim/base.py

@@ -57,11 +57,6 @@ def get_gradients(
     -------
     list of TensorVariable
         One gradient per parameter, in the order of ``parameters``.
-
-    Raises
-    ------
-    DisconnectedInputError
-        If the loss has no gradient with respect to one of ``parameters``, naming the ones it cannot reach.
     """
     if isinstance(loss_or_gradients, list | tuple):
         gradients = list(loss_or_gradients)

pytensor_ml/pytensorf/collect.py

@@ -110,12 +110,6 @@ def collect_clock_updates(
     -------
     clock_updates : dict
         Mapping from each clock the graph reads to the expression for its next value.
-
-    Raises
-    ------
-    ValueError
-        If two clocks the graph reads hold different step counts. They all count training steps, so a
-        disagreement means some of them were restored from a checkpoint and others were not.
     """
     counters = collect_step_counters(outputs)
     step_counts = {int(counter.get_value()) for counter in counters}

pytensor_ml/pytensorf/compile.py

@@ -1,10 +1,11 @@
+import warnings
+
 from collections.abc import Sequence
-from typing import cast
 
 import pytensor
 
 from pytensor import Mode
-from pytensor.compile import Function, SharedVariable, get_mode
+from pytensor.compile import Function, get_mode
 from pytensor.tensor.variable import Variable
 
 from pytensor_ml.pytensorf.collect import collect_graph_inputs
@@ -13,6 +14,8 @@
     SeedSequenceSeed,
     atleast_list,
     collect_default_updates,
+    find_generators_drawn_from,
+    find_rng_nodes,
     reseed_rngs,
 )
 
@@ -37,8 +40,9 @@ def function(
     outputs : Variable or list of Variable
         Outputs of the compiled function.
     random_seed : int, array-like of int, or SeedSequence, optional
-        Seed used to reseed the graph's shared generators. They are replaced whether or not a seed is
-        given, so omitting it reseeds from fresh entropy rather than leaving them untouched.
+        Seed used to replace the graph's shared generators, making the compiled function's draws
+        reproducible. The generators are left as they are when omitted, so compiling has no effect on a
+        generator the caller seeded, or on one another function is drawing from.
     mode : Mode or str, optional
         PyTensor mode used to compile the function.
     **kwargs
@@ -49,16 +53,39 @@ def function(
     Function
         The compiled function.
     """
-    rng_updates = collect_default_updates(
-        inputs=[inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs],
-        outputs=[
-            out.variable if isinstance(out, pytensor.Out) else out for out in atleast_list(outputs)
-        ],
-    )
-
-    if rng_updates:
-        rngs = cast(list[SharedVariable], list(rng_updates))
-        reseed_rngs(rngs, random_seed)
+    updates = dict(kwargs.pop("updates", {}))
+    input_variables = [inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs]
+    # Updates count as readers: a generator a rule draws its noise from is read by the update expression
+    # and often by nothing else, and one read by both an output and an update is read twice.
+    read_variables = [
+        *(out.variable if isinstance(out, pytensor.Out) else out for out in atleast_list(outputs)),
+        *updates.values(),
+    ]
+
+    if random_seed is not None:
+        reseed_rngs(find_rng_nodes(read_variables), random_seed)
+
+    with warnings.catch_warnings():
+        # This warns for a generator with several distinct draws and returns no update for it. The check
+        # below reports that better, and only once the caller's own updates are known.
+        warnings.filterwarnings(
+            "ignore", message="RNG Variable .* multiple distinct clients", category=UserWarning
+        )
+        rng_updates = collect_default_updates(inputs=input_variables, outputs=read_variables)
+
+    frozen = [
+        generator
+        for generator in find_generators_drawn_from(read_variables)
+        if generator not in rng_updates and generator not in updates
+    ]
+    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 "
+            "generator, thread one through with `next_rng, draw = pt.random.normal(rng=rng, "
+            "return_next_rng=True)`, or pass an update for it yourself."
+        )
 
     base_mode = get_mode(mode)
     mode = Mode(
@@ -69,7 +96,7 @@ def function(
     return pytensor.function(
         inputs,
         outputs,
-        updates={**rng_updates, **kwargs.pop("updates", {})},
+        updates={**rng_updates, **updates},
         mode=mode,
         **kwargs,
     )

pytensor_ml/pytensorf/rng.py

@@ -46,6 +46,34 @@ def reseed_rngs(rngs: Sequence[SharedVariable], seed: SeedSequenceSeed) -> None:
         rng.set_value(np.random.Generator(bit_generator), borrow=True)
 
 
+def find_generators_drawn_from(
+    outputs: Sequence[Variable],
+) -> list[RandomGeneratorSharedVariable]:
+    """
+    Return the shared generators the graph draws from, excluding any it only hands back.
+
+    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.
+
+    Parameters
+    ----------
+    outputs : sequence of Variable
+        Graph outputs to trace back from.
+
+    Returns
+    -------
+    list of RandomGeneratorSharedVariable
+        The generators some node in the graph draws from, 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])
+    ]
+
+
 def collect_default_updates_inner_fgraph(node: Apply) -> dict[Variable, Variable]:
     """Collect default RNG updates from a node carrying an inner function graph, mapped to outer variables."""
     op = node.op

tests/test_pytensorf.py

@@ -11,11 +11,13 @@
     zero_grad,
 )
 from pytensor.graph.traversal import ancestors
+from pytensor.tensor import random as ptr
 
 from pytensor_ml.layers import Dropout, DropoutLayer, Linear, Sequential
 from pytensor_ml.pytensorf import (
     collect_trainable_params,
     compile_predict,
+    function,
     rewrite_for_prediction,
     rewrite_pregrad,
 )
@@ -90,3 +92,145 @@ def test_rewrite_pregrad_leaves_a_fully_detached_parameter_disconnected():
 
     with pytest.raises(DisconnectedInputError):
         grad(rewrite_pregrad(loss), W)
+
+
+def test_compiling_leaves_a_generator_another_function_is_drawing_from_alone():
+    """Compiling used to replace every generator it touched, so building a second function mid-training
+    jumped the first one's noise stream."""
+
+    def four_draws(with_intervening_compile):
+        rng = shared(np.random.default_rng(0), name="rng")
+        _, noise = ptr.normal(rng=rng, return_next_rng=True)
+        draw = function([], noise)
+        drawn = [float(draw()) for _ in range(2)]
+        if with_intervening_compile:
+            _, other_noise = ptr.normal(rng=rng, return_next_rng=True)
+            function([], other_noise)
+        return drawn + [float(draw()) for _ in range(2)]
+
+    assert four_draws(with_intervening_compile=True) == four_draws(with_intervening_compile=False)
+
+
+@pytest.mark.parametrize("applications", [1, 2], ids=["applied_once", "applied_twice"])
+def test_a_seeded_dropout_reproduces_across_identical_runs(applications):
+    """The seed a caller puts on a layer has to survive compilation, or `random_state=` means nothing. It has
+    to hold however many times the layer is applied, since each application draws off its own generator."""
+
+    def dropout_masks():
+        X = pt.tensor("X", shape=(None, 3))
+        dropout = Dropout(p=0.5, random_state=0)
+        layers = [Linear("fc", n_in=3, n_out=3)]
+        for _ in range(applications):
+            layers.append(dropout)
+        prediction = Sequential(*layers)(X)
+        for parameter in collect_trainable_params(prediction):
+            parameter.set_value(np.ones_like(parameter.get_value()))
+        forward = function([X], prediction)
+        features = np.ones((4, 3), dtype=config.floatX)
+        return [forward(features) for _ in range(3)]
+
+    for first, second in zip(dropout_masks(), dropout_masks()):
+        np.testing.assert_allclose(first, second)
+
+
+def test_random_seed_makes_an_unseeded_graph_reproducible():
+    def two_draws():
+        rng = shared(np.random.default_rng(), name="rng")  # deliberately unseeded
+        _, noise = ptr.normal(rng=rng, return_next_rng=True)
+        draw = function([], noise, random_seed=42)
+        return [float(draw()) for _ in range(2)]
+
+    assert two_draws() == two_draws()
+
+
+def test_a_reused_dropout_instance_keeps_drawing_new_masks():
+    """Using one Dropout object at two points in a network is an ordinary thing to write, and it used to
+    freeze the mask for the whole run."""
+    X = pt.tensor("X", shape=(None, 3))
+    dropout = Dropout(p=0.5, random_state=0)
+    prediction = Sequential(
+        Linear("fc1", n_in=3, n_out=4), dropout, Linear("fc2", n_in=4, n_out=2), dropout
+    )(X)
+    for parameter in collect_trainable_params(prediction):
+        parameter.set_value(np.ones_like(parameter.get_value()))
+
+    forward = function([X], prediction)
+    features = np.ones((4, 3), dtype=config.floatX)
+    outputs = [forward(features) for _ in range(4)]
+
+    assert (
+        len(dropout.generators) == 2
+    )  # one per application, so neither draw is starved of updates
+    assert all(
+        np.any(outputs[i] != outputs[j]) for i in range(4) for j in range(i + 1, 4)
+    )  # a fresh mask on every call, not just eventually
+
+
+def test_a_returned_generator_is_the_one_the_caller_asked_for():
+    """Returning a generator alongside a draw from it is not two draws: an output reads the generator without
+    consuming it, so it must keep reading the caller's own."""
+    rng = shared(np.random.default_rng(0), name="rng")
+    _, noise = ptr.normal(rng=rng, return_next_rng=True)
+    draw = function([], [noise, rng])
+
+    expected = rng.get_value(borrow=True).bit_generator.state
+    returned = draw()[1].bit_generator.state
+
+    assert returned == expected
+
+
+def test_a_generator_read_only_by_an_update_still_advances():
+    """A rule that adds noise to its step reads a generator the outputs never touch, so collecting updates
+    from the outputs alone left that generator frozen and every step took the identical perturbation."""
+    rng = shared(np.random.default_rng(0), name="rng")
+    parameter = shared(np.zeros(()), name="w")
+    _, noise = ptr.normal(rng=rng, return_next_rng=True)
+
+    step = function([], parameter**2, updates={parameter: parameter + noise})
+
+    increments = []
+    for _ in range(3):
+        before = float(parameter.get_value())
+        step()
+        increments.append(float(parameter.get_value()) - before)
+
+    assert not np.allclose(increments[0], increments[1])
+    assert not np.allclose(increments[1], increments[2])
+
+
+def test_two_distinct_draws_off_one_generator_are_rejected():
+    """One generator cannot serve two different draws: it has no single next state, so nothing can advance
+    it and every call would repeat. pytensor calls such a graph inconsistent and threads no update, which
+    would otherwise show up as a training run whose noise never changes."""
+    rng = shared(np.random.default_rng(0), name="rng")
+    _, normal_draw = ptr.normal(rng=rng, return_next_rng=True)
+    _, uniform_draw = ptr.uniform(rng=rng, return_next_rng=True)
+
+    with pytest.raises(ValueError, match="which nothing advances"):
+        function([], [normal_draw, uniform_draw])
+
+
+def test_a_shared_generator_is_accepted_when_the_caller_advances_it():
+    """The check is about the assembled updates, not about how the graph looks: a caller who threads the
+    generator themselves has answered the question, and their update stands."""
+    rng = shared(np.random.default_rng(0), name="rng")
+    next_rng, normal_draw = ptr.normal(rng=rng, return_next_rng=True)
+    _, uniform_draw = ptr.uniform(rng=next_rng, return_next_rng=True)
+
+    draw = function([], [normal_draw, uniform_draw], updates={rng: next_rng})
+
+    first, second = draw(), draw()
+    assert float(first[0]) != float(second[0])  # the draw off the generator advances
+    assert float(first[1]) != float(second[1])  # and so does the one off the threaded next state
+
+
+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."""
+    rng = shared(np.random.default_rng(0), name="rng")
+    accumulator = shared(np.zeros(()), name="accumulator")
+    _, output_draw = ptr.normal(rng=rng, return_next_rng=True)
+    _, update_draw = ptr.uniform(rng=rng, return_next_rng=True)
+
+    with pytest.raises(ValueError, match="which nothing advances"):
+        function([], output_draw, updates={accumulator: accumulator + update_draw})