Let compile_train take extra updates

pytensor_ml/optim/train.py

@@ -1,10 +1,10 @@
 from collections.abc import Sequence
 
-from pytensor.compile import Function, SharedVariable
+from pytensor.compile import Function
 from pytensor.graph.basic import Variable, equal_computations
 from pytensor.tensor import TensorVariable
 
-from pytensor_ml.optim.base import Parameter, UpdateRule, require_unique_state_names
+from pytensor_ml.optim.base import Parameter, UpdateRule, Updates, require_unique_state_names
 from pytensor_ml.pytensorf import (
     collect_clock_updates,
     collect_data_inputs,
@@ -21,14 +21,16 @@ def compile_train(
     parameters: Sequence[Parameter] | None = None,
     inputs: Sequence[Variable] | None = None,
     extra_outputs: Sequence[Variable] | None = None,
+    extra_updates: Updates | None = None,
     compile_kwargs: dict | None = None,
 ) -> Function:
     """
     Compile a one-step training function from a loss graph and an update rule.
 
     Differentiates the loss via ``rule``, applies the resulting updates, folds in any non-trainable state
-    updates (such as batch-norm running statistics), advances every training clock the step reads, and
-    compiles. The parameters and data inputs are collected from ``loss`` unless given explicitly.
+    updates (such as batch-norm running statistics) and any given in ``extra_updates``, advances every
+    training clock the step reads, and compiles. The parameters and data inputs are collected from ``loss``
+    unless given explicitly.
 
     Parameters
     ----------
@@ -42,17 +44,25 @@ def compile_train(
         is still initialized and checkpointed, since :func:`collect_trainable_params` reaches it; only the
         optimizer skips it.
     inputs : sequence of Variable, optional
-        Data inputs of the compiled function, in call order. Collected from ``loss`` and ``extra_outputs``
-        with :func:`collect_data_inputs` when omitted; pass them explicitly when call order matters (e.g.
-        features before targets).
+        Data inputs of the compiled function, in call order. Collected from ``loss``, ``extra_outputs`` and
+        ``extra_updates`` with :func:`collect_data_inputs` when omitted; pass them explicitly when call order
+        matters (e.g. features before targets).
     extra_outputs : sequence of Variable, optional
         Diagnostics to return alongside the loss, such as gradient norms or a batch accuracy. Evaluated in
         the same pass as the gradients, so they see the pre-update parameter values, and they add no
         non-trainable state updates of their own. A random node reached only through an extra output does
         still advance its generator, since :func:`~pytensor_ml.pytensorf.function` threads the next-RNG
         update for every generator the outputs draw from.
+    extra_updates : dict, optional
+        State the step should write that no gradient produces -- a target-network sync, a Polyak average,
+        replay priorities. Mapping from shared variable to its next value, folded in alongside the rule's own
+        updates. Raise if the rule or the model already writes one of these variables, since two writes to one
+        variable cannot both take effect. An expression here is part of the step like any other, so a clock it
+        reads still advances once, and a generator it draws from still advances.
     compile_kwargs : dict, optional
-        Extra keyword arguments forwarded to the function compiler.
+        Extra keyword arguments forwarded to the function compiler. An ``updates`` entry is taken as
+        ``extra_updates`` rather than forwarded, since the compiler's own ``updates`` carries the whole
+        assembled step. Raise if a variable is given in both.
 
     Returns
     -------
@@ -61,19 +71,45 @@ def compile_train(
         ``(loss, *extra_outputs)`` when diagnostics were requested.
     """
     extra_outputs = list(extra_outputs or [])
+    extra_updates = dict(extra_updates or {})
+    # Copied rather than mutated, so popping the caller's `updates` out does not empty their dict.
+    compile_kwargs = dict(compile_kwargs or {})
+
+    # `updates` is pytensor's own name for this, so a caller who puts one in compile_kwargs is asking for
+    # what extra_updates does; take it as one instead of letting it collide with this function's own call.
+    for variable, new_value in compile_kwargs.pop("updates", {}).items():
+        if variable in extra_updates:
+            raise ValueError(
+                f"The update for {variable.name!r} is given twice, once in `extra_updates` and once in "
+                "`compile_kwargs['updates']`. Both are folded into the training step, so keep whichever one "
+                "is right and drop the other."
+            )
+        extra_updates[variable] = new_value
 
     if parameters is None:
         parameters = collect_differentiable_params(loss)
     if inputs is None:
-        inputs = collect_data_inputs([loss, *extra_outputs])
+        inputs = collect_data_inputs([loss, *extra_outputs, *extra_updates.values()])
 
-    updates: dict[SharedVariable, TensorVariable] = dict(rule(loss, parameters))
+    updates: Updates = dict(rule(loss, parameters))
 
     # Assigned per key rather than merged: SupportsKeysAndGetItem is invariant in its key type, so
     # dict.update rejects the narrower NonTrainableParameter keys.
     for parameter, new_value in collect_non_trainable_updates(loss).items():
         updates[parameter] = new_value
 
+    # Folded in after the rule's and the model's, so a collision with either is caught rather than deciding
+    # by insertion order which of the two writes survives.
+    for variable, new_value in extra_updates.items():
+        if variable in updates:
+            raise ValueError(
+                f"The extra update for {variable.name!r} writes a variable the training step already "
+                "writes, so the two writes cannot both take effect. Optimizer state and batch-norm "
+                "statistics are written by the step itself; drop this one, or fold what it does into the "
+                "expression that already writes the variable."
+            )
+        updates[variable] = new_value
+
     # Collected from the assembled updates rather than from the loss alone: a clock is read by a schedule
     # or a policy, which live in the updates, where an RNG or a running statistic is read by the model.
     for clock, next_count in collect_clock_updates(
@@ -96,4 +132,4 @@ def compile_train(
 
     outputs = [loss, *extra_outputs] if extra_outputs else loss
 
-    return function(list(inputs), outputs, updates=updates, **(compile_kwargs or {}))
+    return function(list(inputs), outputs, updates=updates, **compile_kwargs)

tests/optim/test_training.py

@@ -2,17 +2,17 @@
 import pytensor.tensor as pt
 import pytest
 
-from pytensor import config
+from pytensor import config, shared
 from pytensor.gradient import disconnected_grad, zero_grad
 from sklearn.datasets import load_digits, make_regression
 from sklearn.preprocessing import MinMaxScaler, OneHotEncoder, StandardScaler
 
 from pytensor_ml.activations import LeakyReLU
 from pytensor_ml.layers import BatchNorm2D, Linear, Sequential
 from pytensor_ml.loss import CrossEntropy, SquaredError, supervised_loss
-from pytensor_ml.optim import adam, adamw, compile_train, sgd
+from pytensor_ml.optim import adam, adamw, compile_train, cosine_schedule, sgd
 from pytensor_ml.optim.base import state_for
-from pytensor_ml.params import trainable
+from pytensor_ml.params import step_counter, trainable
 from pytensor_ml.pytensorf import collect_non_trainable_params, collect_trainable_params
 from pytensor_ml.state import initialize_params
 from pytensor_ml.util import DataLoader
@@ -307,3 +307,169 @@ def test_compile_train_leaves_a_zero_grad_parameter_untouched_under_weight_decay
         for parameter, value in before.items()
     }
     assert moved == {"live_W": True, "live_b": True, "frozen_W": False, "frozen_b": False}
+
+
+def test_extra_updates_write_state_no_gradient_produces():
+    # The DQN shape from the other side: a target network kept as a Polyak average of the online weights.
+    # No gradient produces that write, so without extra_updates it cannot ride along in the training step.
+    X = pt.tensor("X", shape=(None, 4))
+    online_layer = Linear("online", n_in=4, n_out=2)
+    prediction = online_layer(X)
+    parameters = collect_trainable_params(prediction)
+    initialize(parameters)
+    target_weight = shared(online_layer.W.get_value().copy(), name="target_W")
+    loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2)
+
+    step = compile_train(
+        loss,
+        sgd(1e-1),
+        parameters=parameters,
+        inputs=[X, target],
+        extra_updates={target_weight: 0.5 * target_weight + 0.5 * online_layer.W},
+    )
+
+    online_before = online_layer.W.get_value().copy()
+    target_before = target_weight.get_value().copy()
+
+    rng = np.random.default_rng(0)
+    features = rng.normal(size=(16, 4)).astype(config.floatX)
+    targets = np.ones((16, 2), dtype=config.floatX)
+    step(features, targets)
+
+    # Updates are computed from the pre-update values, so the target lands halfway between where the two
+    # started -- neither staying put nor jumping to where the online weights ended up.
+    assert not np.allclose(online_layer.W.get_value(), online_before)
+    np.testing.assert_allclose(
+        target_weight.get_value(), 0.5 * target_before + 0.5 * online_before, rtol=1e-5
+    )
+
+
+def test_extra_updates_reject_a_write_the_rule_already_makes():
+    # Silently overwriting an optimizer buffer would leave the rule configured but not working, for the whole
+    # run, so the collision has to be loud.
+    p = trainable(np.array([2.0]), name="w")
+    loss = 0.5 * (p**2).sum()
+    rule = adam(1e-1)
+    first_moment = next(key for key in rule(loss, [p]) if key.name == "w/adam/first_moment")
+
+    with pytest.raises(ValueError, match="already writes"):
+        compile_train(loss, rule, extra_updates={first_moment: first_moment * 0.0}, inputs=[])
+
+
+def test_extra_updates_reject_a_write_the_model_already_makes():
+    # Batch-norm statistics are written by the model rather than the rule, and collide just the same.
+    X = pt.tensor("X", shape=(None, 4))
+    prediction = Sequential(Linear("fc", n_in=4, n_out=4), BatchNorm2D("bn", n_in=4))(X)
+    parameters = collect_trainable_params(prediction)
+    initialize(parameters)
+    loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2)
+    running_mean = next(
+        p for p in collect_non_trainable_params(prediction) if "running_mean" in p.name
+    )
+
+    with pytest.raises(ValueError, match="already writes"):
+        compile_train(
+            loss,
+            sgd(1e-2),
+            parameters=parameters,
+            inputs=[X, target],
+            extra_updates={running_mean: running_mean * 0.0},
+        )
+
+
+def test_extra_updates_contribute_their_data_inputs():
+    # An extra update may read data the loss never touches -- replay priorities, an importance weight. Unless
+    # those inputs are collected too, compiling raises MissingInputError.
+    p = trainable(np.array([2.0]), name="w")
+    priorities = shared(np.zeros(3, dtype=config.floatX), name="priorities")
+    fresh_priorities = pt.vector("fresh_priorities", shape=(3,))
+    loss = 0.5 * (p**2).sum()
+
+    step = compile_train(loss, sgd(1e-1), extra_updates={priorities: fresh_priorities})
+
+    step(fresh_priorities=np.array([1.0, 2.0, 3.0], dtype=config.floatX))
+    np.testing.assert_allclose(priorities.get_value(), [1.0, 2.0, 3.0])
+
+
+def test_an_extra_update_reads_the_clock_and_advances_it_once():
+    # An extra update is part of the step, so a clock it reads is a clock the step reads: it advances once
+    # per step, and the expression sees the count from before the advance, like the rule's updates do.
+    p = trainable(np.array([2.0]), name="w")
+    decayed = shared(np.array(1.0, dtype=config.floatX), name="decayed")
+    clock = step_counter(name="training_step")
+    schedule = cosine_schedule(1.0, 10)
+    loss = 0.5 * (p**2).sum()
+
+    step = compile_train(loss, sgd(1e-1), extra_updates={decayed: schedule(clock)}, inputs=[])
+    for _ in range(3):
+        step()
+
+    assert int(clock.get_value()) == 3
+    expected = float(schedule(pt.as_tensor(2, dtype="int64")).eval())
+    np.testing.assert_allclose(decayed.get_value(), expected, rtol=1e-6)
+
+
+def test_an_extra_update_that_draws_noise_advances_its_generator():
+    # A perturbed step -- SGLD, exploration noise -- draws inside the update rather than the loss. Nothing
+    # else reads that generator, so unless the step advances it every call adds the identical perturbation.
+    p = trainable(np.array([2.0]), name="w")
+    perturbed = shared(np.zeros(3, dtype=config.floatX), name="perturbed")
+    noise_rng = shared(np.random.default_rng(0), name="noise_rng")
+    _, noise = pt.random.normal(size=(3,), rng=noise_rng, return_next_rng=True)
+    loss = 0.5 * (p**2).sum()
+
+    step = compile_train(
+        loss, sgd(1e-1), extra_updates={perturbed: noise.astype(config.floatX)}, inputs=[]
+    )
+    step()
+    first = perturbed.get_value().copy()
+    step()
+
+    assert not np.allclose(first, perturbed.get_value())
+
+
+def test_updates_in_compile_kwargs_are_taken_as_extra_updates():
+    # `updates` is what pytensor calls this, so it is the first place a caller looks. Forwarding it would
+    # collide with the compiler's own updates argument and raise a TypeError naming an internal function.
+    p = trainable(np.array([2.0]), name="w")
+    call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count")
+    loss = 0.5 * (p**2).sum()
+
+    step = compile_train(
+        loss, sgd(1e-1), inputs=[], compile_kwargs={"updates": {call_count: call_count + 1.0}}
+    )
+    step()
+    step()
+
+    assert float(call_count.get_value()) == 2.0
+
+
+def test_an_update_given_in_both_places_is_rejected():
+    p = trainable(np.array([2.0]), name="w")
+    call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count")
+    loss = 0.5 * (p**2).sum()
+
+    with pytest.raises(ValueError, match="given twice"):
+        compile_train(
+            loss,
+            sgd(1e-1),
+            inputs=[],
+            extra_updates={call_count: call_count + 1.0},
+            compile_kwargs={"updates": {call_count: call_count + 2.0}},
+        )
+
+
+def test_compile_kwargs_is_not_mutated():
+    # Taking `updates` out of the caller's dict would empty it, so compiling twice from one settings dict
+    # would silently drop the update the second time.
+    p = trainable(np.array([2.0]), name="w")
+    call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count")
+    loss = 0.5 * (p**2).sum()
+    compile_kwargs = {"updates": {call_count: call_count + 1.0}}
+
+    compile_train(loss, sgd(1e-1), inputs=[], compile_kwargs=compile_kwargs)
+    second = compile_train(loss, sgd(1e-1), inputs=[], compile_kwargs=compile_kwargs)
+    second()
+
+    assert "updates" in compile_kwargs
+    assert float(call_count.get_value()) == 1.0