From e1df8a39e7820df1e9c966b232e83428bf426338 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:25:18 +0800
Subject: [PATCH 01/16] Accept a schedule wherever a learning rate goes

---
 pytensor_ml/optim/alias.py     | 185 ++++++++++++---------------------
 pytensor_ml/optim/base.py      |  53 +++++++++-
 pytensor_ml/optim/rules.py     |  56 +++++-----
 pytensor_ml/optim/transform.py |  39 ++++---
 tests/optim/test_rules.py      |  16 +++
 tests/optim/test_transform.py  |  15 +++
 tests/test_checkpoint.py       |  15 ++-
 7 files changed, 209 insertions(+), 170 deletions(-)

diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py
index 17d2210..4f9c3dd 100644
--- a/pytensor_ml/optim/alias.py
+++ b/pytensor_ml/optim/alias.py
@@ -4,10 +4,8 @@
     LearningRate,
     LossGradientsOrUpdates,
     Parameter,
-    Rate,
     Transform,
     Updates,
-    counter,
     reuses_state,
 )
 from pytensor_ml.optim.rules import (
@@ -25,20 +23,6 @@
 from pytensor_ml.optim.transform import scale, trace
 
 
-def _at_learning_rate(
-    learning_rate: LearningRate,
-    name: str,
-    build_updates: Callable[[Rate], Updates],
-) -> Updates:
-    """Build updates at ``learning_rate``, reading a schedule off the clock the rule counts its own steps on.
-    Both reach that clock through :func:`counter` under ``"{name}/step_count"``, so a rule that keeps a step
-    count hands the schedule the same variable instead of a second one measuring the same time."""
-    if not callable(learning_rate):
-        return build_updates(learning_rate)
-
-    return build_updates(learning_rate(counter(f"{name}/step_count")))
-
-
 def sgd(
     learning_rate: LearningRate = 0.01,
     momentum: float = 0.0,
@@ -94,24 +78,21 @@ def sgd(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        def build_updates(rate: Rate) -> Updates:
-            if momentum_trace is None:
-                return sgd_updates(
-                    loss_gradients_or_updates,
-                    parameters,
-                    learning_rate=rate,
-                    namespace=namespace,
-                )
-            updates = sgd_updates(
+        if momentum_trace is None:
+            return sgd_updates(
                 loss_gradients_or_updates,
                 parameters,
-                learning_rate=1.0,
+                learning_rate=learning_rate,
                 namespace=namespace,
             )
-            updates = momentum_trace(updates, parameters)
-            return scale(rate)(updates, parameters)
-
-        return _at_learning_rate(learning_rate, namespace, build_updates)
+        updates = sgd_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=1.0,
+            namespace=namespace,
+        )
+        updates = momentum_trace(updates, parameters)
+        return scale(learning_rate, namespace=namespace)(updates, parameters)
 
     return rule
 
@@ -155,19 +136,15 @@ def adam(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: adam_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                beta1=beta1,
-                beta2=beta2,
-                epsilon=epsilon,
-                amsgrad=amsgrad,
-                namespace=namespace,
-            ),
+        return adam_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            beta1=beta1,
+            beta2=beta2,
+            epsilon=epsilon,
+            amsgrad=amsgrad,
+            namespace=namespace,
         )
 
     return rule
@@ -215,21 +192,17 @@ def adamw(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: adamw_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                weight_decay=weight_decay,
-                beta1=beta1,
-                beta2=beta2,
-                epsilon=epsilon,
-                amsgrad=amsgrad,
-                mask=mask,
-                namespace=namespace,
-            ),
+        return adamw_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            weight_decay=weight_decay,
+            beta1=beta1,
+            beta2=beta2,
+            epsilon=epsilon,
+            amsgrad=amsgrad,
+            mask=mask,
+            namespace=namespace,
         )
 
     return rule
@@ -274,18 +247,14 @@ def nadam(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: nadam_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                beta1=beta1,
-                beta2=beta2,
-                epsilon=epsilon,
-                namespace=namespace,
-            ),
+        return nadam_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            beta1=beta1,
+            beta2=beta2,
+            epsilon=epsilon,
+            namespace=namespace,
         )
 
     return rule
@@ -330,18 +299,14 @@ def adamax(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: adamax_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                beta1=beta1,
-                beta2=beta2,
-                epsilon=epsilon,
-                namespace=namespace,
-            ),
+        return adamax_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            beta1=beta1,
+            beta2=beta2,
+            epsilon=epsilon,
+            namespace=namespace,
         )
 
     return rule
@@ -442,19 +407,15 @@ def rmsprop(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: rmsprop_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                rho=rho,
-                momentum=momentum,
-                epsilon=epsilon,
-                centered=centered,
-                namespace=namespace,
-            ),
+        return rmsprop_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            rho=rho,
+            momentum=momentum,
+            epsilon=epsilon,
+            centered=centered,
+            namespace=namespace,
         )
 
     return rule
@@ -496,16 +457,12 @@ def adagrad(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: adagrad_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                epsilon=epsilon,
-                namespace=namespace,
-            ),
+        return adagrad_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            epsilon=epsilon,
+            namespace=namespace,
         )
 
     return rule
@@ -548,17 +505,13 @@ def adadelta(
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        return _at_learning_rate(
-            learning_rate,
-            namespace,
-            lambda rate: adadelta_updates(
-                loss_gradients_or_updates,
-                parameters,
-                learning_rate=rate,
-                rho=rho,
-                epsilon=epsilon,
-                namespace=namespace,
-            ),
+        return adadelta_updates(
+            loss_gradients_or_updates,
+            parameters,
+            learning_rate=learning_rate,
+            rho=rho,
+            epsilon=epsilon,
+            namespace=namespace,
         )
 
     return rule
diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 07f3a96..62b840b 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -179,12 +179,12 @@ def inverse_square_root(step_count):
 
 type LearningRate = Rate | Schedule
 """
-What an optimizer alias accepts as its rate, adding a schedule that drives it on-graph.
+What any ``learning_rate`` accepts: a rate, or a schedule the rule reads off its own training clock.
 
 Examples
 --------
-Every alias takes either form, so a constant can be swapped for a schedule without touching anything
-else:
+Every rule and alias takes either form, so a constant can be swapped for a schedule without touching
+anything else:
 
 .. code-block:: python
 
@@ -195,6 +195,53 @@ def inverse_square_root(step_count):
 """
 
 
+def rate_on(learning_rate: LearningRate, clock: Parameter) -> Rate:
+    """
+    Read a schedule off ``clock``. Any other rate passes through untouched.
+
+    Parameters
+    ----------
+    learning_rate : LearningRate
+        A rate, or a schedule of the step count.
+    clock : shared tensor variable
+        The training clock a schedule is evaluated at, ordinarily the one the rule counts its own steps on.
+
+    Returns
+    -------
+    rate : Rate
+        The rate as a number or a scalar graph.
+    """
+    return learning_rate(clock) if callable(learning_rate) else learning_rate
+
+
+def read_rate(learning_rate: LearningRate, namespace: str) -> tuple[Rate, Updates]:
+    """
+    Resolve a rate at ``floatX``, reading a schedule off a clock of the caller's own.
+
+    The clock's advance comes back as an update for the caller to write, so the updates dict carries the
+    clock and a checkpoint taken from the dict resumes the schedule where it left off. A plain rate reads
+    no clock, so none is allocated and nothing comes back to write.
+
+    Parameters
+    ----------
+    learning_rate : LearningRate
+        A rate, or a schedule of the step count.
+    namespace : str
+        Prefix of the clock a schedule reads, allocated as ``"{namespace}/step_count"``.
+
+    Returns
+    -------
+    rate : Rate
+        The rate as a number or a scalar graph.
+    clock_update : Updates
+        The clock's one-step advance when a schedule was read, empty otherwise.
+    """
+    if not callable(learning_rate):
+        return to_floatx(learning_rate), Updates()
+    clock = counter(f"{namespace}/step_count")
+    return to_floatx(learning_rate(clock)), Updates({clock: clock + 1})
+
+
 def to_floatx(value: Rate) -> Rate:
     """
     Return ``value`` at the current ``floatX``, casting only a variable stored at something else.
diff --git a/pytensor_ml/optim/rules.py b/pytensor_ml/optim/rules.py
index bb5b450..e10f80b 100644
--- a/pytensor_ml/optim/rules.py
+++ b/pytensor_ml/optim/rules.py
@@ -3,6 +3,7 @@
 import pytensor.tensor as pt
 
 from pytensor import config
+from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph.basic import Variable
 from pytensor.tensor import TensorVariable
 
@@ -10,11 +11,12 @@
     LearningRate,
     LossGradientsOrUpdates,
     Parameter,
-    Rate,
     Steps,
     Updates,
     counter,
     gradients_to_descend,
+    rate_on,
+    read_rate,
     state_for,
     to_floatx,
 )
@@ -23,7 +25,7 @@
 def sgd_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 1.0,
+    learning_rate: LearningRate = 1.0,
     namespace: str = "sgd",
 ) -> Updates:
     r"""
@@ -73,19 +75,18 @@ def sgd_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
-    return Steps(incoming).replacing(
-        {
-            parameter: parameter - learning_rate * gradient
-            for parameter, gradient in zip(parameters, gradients)
-        }
-    )
+    learning_rate, clock_update = read_rate(learning_rate, namespace)
+    steps: dict[SharedVariable, TensorVariable] = {
+        parameter: parameter - learning_rate * gradient
+        for parameter, gradient in zip(parameters, gradients)
+    }
+    return Steps(incoming).replacing(clock_update).replacing(steps)
 
 
 def adam_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 1e-3,
+    learning_rate: LearningRate = 1e-3,
     beta1: float = 0.9,
     beta2: float = 0.999,
     epsilon: float = 1e-8,
@@ -166,7 +167,7 @@ def _adam_family_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
     *,
-    learning_rate: Rate,
+    learning_rate: LearningRate,
     beta1: float,
     beta2: float,
     epsilon: float,
@@ -189,9 +190,8 @@ def _adam_family_updates(
         Predicate selecting which parameters the decay reaches. Every parameter when omitted.
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
-
     step_count = counter(f"{namespace}/step_count")
+    learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
     first_moment_bias_correction = 1 - beta1**new_step_count_float
@@ -254,7 +254,7 @@ def _running_max_second_moment(
 def adamw_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 1e-3,
+    learning_rate: LearningRate = 1e-3,
     weight_decay: float = 0.01,
     beta1: float = 0.9,
     beta2: float = 0.999,
@@ -344,7 +344,7 @@ def adamw_updates(
 def nadam_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 2e-3,
+    learning_rate: LearningRate = 2e-3,
     beta1: float = 0.9,
     beta2: float = 0.999,
     epsilon: float = 1e-8,
@@ -408,9 +408,8 @@ def nadam_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
-
     step_count = counter(f"{namespace}/step_count")
+    learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
     first_moment_bias_correction = 1 - beta1**new_step_count_float
@@ -442,7 +441,7 @@ def nadam_updates(
 def adamax_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 2e-3,
+    learning_rate: LearningRate = 2e-3,
     beta1: float = 0.9,
     beta2: float = 0.999,
     epsilon: float = 1e-8,
@@ -505,9 +504,8 @@ def adamax_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
-
     step_count = counter(f"{namespace}/step_count")
+    learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
     first_moment_bias_correction = 1 - beta1**new_step_count_float
@@ -533,7 +531,7 @@ def adamax_updates(
 def adagrad_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 0.01,
+    learning_rate: LearningRate = 0.01,
     epsilon: float = 1e-8,
     namespace: str = "adagrad",
 ) -> Updates:
@@ -588,9 +586,9 @@ def adagrad_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
+    learning_rate, clock_update = read_rate(learning_rate, namespace)
 
-    updates: Updates = Steps(incoming)
+    updates: Updates = Steps(incoming).replacing(clock_update)
     for parameter, gradient in zip(parameters, gradients):
         sum_squared_gradients = state_for(parameter, f"{namespace}/sum_squared_gradients")
         new_sum_squared_gradients = sum_squared_gradients + gradient**2
@@ -605,7 +603,7 @@ def adagrad_updates(
 def rmsprop_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 1e-2,
+    learning_rate: LearningRate = 1e-2,
     rho: float = 0.9,
     momentum: float = 0.0,
     epsilon: float = 1e-8,
@@ -673,9 +671,9 @@ def rmsprop_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
+    learning_rate, clock_update = read_rate(learning_rate, namespace)
 
-    updates: Updates = Steps(incoming)
+    updates: Updates = Steps(incoming).replacing(clock_update)
     for parameter, gradient in zip(parameters, gradients):
         mean_squared_gradient = state_for(parameter, f"{namespace}/mean_squared_gradient")
         new_mean_squared_gradient = rho * mean_squared_gradient + (1 - rho) * gradient**2
@@ -704,7 +702,7 @@ def rmsprop_updates(
 def adadelta_updates(
     loss_gradients_or_updates: LossGradientsOrUpdates,
     parameters: Sequence[Parameter],
-    learning_rate: Rate = 1.0,
+    learning_rate: LearningRate = 1.0,
     rho: float = 0.9,
     epsilon: float = 1e-8,
     namespace: str = "adadelta",
@@ -764,9 +762,9 @@ def adadelta_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    learning_rate = to_floatx(learning_rate)
+    learning_rate, clock_update = read_rate(learning_rate, namespace)
 
-    updates: Updates = Steps(incoming)
+    updates: Updates = Steps(incoming).replacing(clock_update)
     for parameter, gradient in zip(parameters, gradients):
         accumulated_squared_gradient = state_for(
             parameter, f"{namespace}/accumulated_squared_gradient"
diff --git a/pytensor_ml/optim/transform.py b/pytensor_ml/optim/transform.py
index 844aa8b..27aebd7 100644
--- a/pytensor_ml/optim/transform.py
+++ b/pytensor_ml/optim/transform.py
@@ -1,16 +1,17 @@
 from collections.abc import Callable, Sequence
 
+from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.tensor import TensorVariable
 
 from pytensor_ml.optim.base import (
     Gradients,
+    LearningRate,
     LossGradientsOrUpdates,
     Parameter,
-    Rate,
     Schedule,
     Transform,
     Updates,
-    counter,
+    read_rate,
     reuses_state,
     state_for,
     steps_of,
@@ -94,9 +95,9 @@ def transform(
     return transform
 
 
-def scale(factor: Rate) -> Transform:
+def scale(factor: LearningRate, *, namespace: str = "scale") -> Transform:
     """
-    Scale each step by a constant factor.
+    Scale each step by a factor, or by a schedule read off a training clock of its own.
 
     Typically the terminal transform in a chain, used to apply the learning rate after a unit-rate base rule.
 
@@ -105,8 +106,11 @@ def scale(factor: Rate) -> Transform:
 
     Parameters
     ----------
-    factor : float or shared tensor variable
-        Multiplier applied to every step.
+    factor : float, shared tensor variable, symbolic scalar, or Schedule
+        Multiplier applied to every step. A schedule is read off the clock ``"{namespace}/step_count"``.
+    namespace : str
+        Prefix for the clock a schedule reads, so two scheduled scalings in one chain keep separate clocks.
+        Default ``"scale"``.
 
     Returns
     -------
@@ -133,17 +137,18 @@ def scale(factor: Rate) -> Transform:
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
+    @reuses_state
     def transform(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
         updates = to_updates(loss_gradients_or_updates, parameters)
         _reject_gradients(updates, f"scale({factor})")
-        return updates.replacing(
-            {
-                parameter: parameter + factor * step
-                for parameter, step in zip(parameters, steps_of(updates, parameters))
-            }
-        )
+        rate, clock_update = read_rate(factor, namespace)
+        scaled: dict[SharedVariable, TensorVariable] = {
+            parameter: parameter + rate * step
+            for parameter, step in zip(parameters, steps_of(updates, parameters))
+        }
+        return updates.replacing(clock_update).replacing(scaled)
 
     return transform
 
@@ -201,15 +206,7 @@ def scale_by_schedule(schedule: Schedule, *, namespace: str = "scale_by_schedule
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
-    def transform(
-        loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
-    ) -> Updates:
-        updates = to_updates(loss_gradients_or_updates, parameters)
-        _reject_gradients(updates, "scale_by_schedule")
-        return scale(schedule(counter(f"{namespace}/step_count")))(updates, parameters)
-
-    return transform
+    return scale(schedule, namespace=namespace)
 
 
 def add_weight_decay(
diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py
index 1afab51..9ed0aa7 100644
--- a/tests/optim/test_rules.py
+++ b/tests/optim/test_rules.py
@@ -623,3 +623,19 @@ def test_get_gradients_names_a_parameter_lost_to_a_second_derivative():
 
     with pytest.raises(DisconnectedInputError, match=r"\['output_bias'\]"):
         sgd_updates(loss, [weight, scale, output_bias])
+
+
+def test_a_functional_rule_reads_a_schedule_off_its_own_clock():
+    """A schedule is a learning rate, so the functional API takes one where it takes a float. The rule
+    resolves it against the clock it already counts its own steps on, so the graph holds one clock."""
+    p = trainable(np.zeros(3), name="w")
+    loss = (p**2).sum()
+    updates = adam_updates(loss, [p], learning_rate=cosine_schedule(0.1, total_steps=10))
+    clocks = [key for key in updates if isinstance(key, params.StepCounter)]
+
+    step = function([], loss, updates=updates)
+    step()
+    step()
+
+    assert [clock.name for clock in clocks] == ["adam/step_count"]
+    assert int(clocks[0].get_value()) == 2
diff --git a/tests/optim/test_transform.py b/tests/optim/test_transform.py
index 7602ecb..746b88b 100644
--- a/tests/optim/test_transform.py
+++ b/tests/optim/test_transform.py
@@ -151,3 +151,18 @@ def test_add_weight_decay_skips_masked_params():
     p = trainable(np.array([4.0]), name="bias")
     out = add_weight_decay(0.1, mask=lambda param: "bias" not in param.name)({p: p}, [p])
     np.testing.assert_allclose(function([], out[p])(), [4.0])
+
+
+def test_scale_takes_a_schedule_where_it_takes_a_factor():
+    """A schedule is a learning rate, so ``scale`` takes one where it takes a float and reads it off a
+    clock of its own, which is what ``scale_by_schedule`` spells out."""
+    p = trainable(np.zeros(2), name="w")
+    updates = {p: p + pt.constant(np.array([2.0, -4.0]))}
+    out = scale(linear_schedule(1.0, total_steps=4, final_learning_rate=0.0))(updates, [p])
+
+    (clock,) = collect_step_counters(out[p])
+    step = function([], out[p], updates={clock: clock.advance()})
+    step()
+
+    assert clock.name == "scale/step_count"
+    np.testing.assert_allclose(step(), [1.5, -3.0])
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
index cc2f2e6..9cdbdf0 100644
--- a/tests/test_checkpoint.py
+++ b/tests/test_checkpoint.py
@@ -12,7 +12,7 @@
 from pytensor_ml.checkpoint import generator_from_state, load_state, save_state
 from pytensor_ml.layers import BatchNorm, Dropout, Linear, Sequential
 from pytensor_ml.loss import SquaredError, supervised_loss
-from pytensor_ml.optim import adam, sgd
+from pytensor_ml.optim import adam, cosine_schedule, sgd
 from pytensor_ml.params import trainable
 from pytensor_ml.pytensorf import (
     collect_optimizer_state,
@@ -631,3 +631,16 @@ def test_collect_optimizer_state_is_empty_for_a_rule_that_keeps_none():
     parameters = collect_trainable_params(loss)
 
     assert collect_optimizer_state(sgd(learning_rate=1e-2)(loss, parameters), parameters) == []
+
+
+def test_a_scheduled_rule_puts_the_clock_it_reads_in_its_updates():
+    """A schedule's clock is state the rule carries between steps, so it has to be in the updates dict a
+    checkpoint is taken from. A rule that read the clock without writing it would train correctly and
+    still hand a checkpoint that restarts the schedule at zero."""
+    X = pt.tensor("X", shape=(None, 4))
+    loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())
+    parameters = collect_trainable_params(loss)
+
+    updates = sgd(learning_rate=cosine_schedule(1e-2, total_steps=10))(loss, parameters)
+
+    assert [v.name for v in collect_optimizer_state(updates, parameters)] == ["sgd/step_count"]

From a1561597be1e0bb598e690bbd872e80d31e2dd9f Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:27:54 +0800
Subject: [PATCH 02/16] Allocate optimizer state afresh on every rule
 invocation

A rule is a graph builder, and the updates dict it returns is the state-bearing object, so two functions that share state are compiled from one dict. The memoizing scope this replaces closed over a ContextVar, which is why a configured rule could not be pickled. Closes #160.
---
 conda_envs/pytensor_ml.yml         |   1 +
 docs/source/api/optim.rst          |   1 -
 pyproject.toml                     |   1 +
 pytensor_ml/optim/__init__.py      |   2 -
 pytensor_ml/optim/alias.py         |  18 +---
 pytensor_ml/optim/base.py          | 160 ++++-------------------------
 pytensor_ml/optim/guards.py        |   2 -
 pytensor_ml/optim/policy.py        |   2 -
 pytensor_ml/optim/train.py         |  11 +-
 pytensor_ml/optim/transform.py     |   3 -
 tests/optim/test_composition.py    |  16 +++
 tests/optim/test_learning_rate.py  |   6 +-
 tests/optim/test_policy.py         |   6 +-
 tests/optim/test_public_surface.py |   4 +-
 tests/optim/test_rules.py          |  38 +++----
 tests/optim/test_training.py       |  19 ++--
 tests/optim/test_transform.py      |  27 -----
 tests/optim/test_update_space.py   |  54 +---------
 18 files changed, 77 insertions(+), 294 deletions(-)

diff --git a/conda_envs/pytensor_ml.yml b/conda_envs/pytensor_ml.yml
index fba4723..36717ef 100644
--- a/conda_envs/pytensor_ml.yml
+++ b/conda_envs/pytensor_ml.yml
@@ -11,6 +11,7 @@ dependencies:
   - scikit-learn
 
   # For testing
+  - cloudpickle  # tests/optim/test_composition.py ships a configured rule to a worker
   - coveralls
   - diff-cover
   - pytest
diff --git a/docs/source/api/optim.rst b/docs/source/api/optim.rst
index 5607bb4..2391691 100644
--- a/docs/source/api/optim.rst
+++ b/docs/source/api/optim.rst
@@ -88,7 +88,6 @@ Building blocks
 
     counter
     get_gradients
-    reuses_state
     scalar_state
     state_for
     steps_of
diff --git a/pyproject.toml b/pyproject.toml
index b731fc2..73fc934 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,6 +48,7 @@ Tracker = "https://github.com/pymc-devs/pytensor-ml/issues"
 [project.optional-dependencies]
 dev = [
     "safetensors",
+    "cloudpickle",
     "pre-commit",
     "pytest",
     "pytest-cov",
diff --git a/pytensor_ml/optim/__init__.py b/pytensor_ml/optim/__init__.py
index d8205b1..5f0a595 100644
--- a/pytensor_ml/optim/__init__.py
+++ b/pytensor_ml/optim/__init__.py
@@ -20,7 +20,6 @@
     chain,
     counter,
     get_gradients,
-    reuses_state,
     scalar_state,
     state_for,
     steps_of,
@@ -106,7 +105,6 @@
     "nonfinite",
     "polynomial_schedule",
     "reduce_on_plateau",
-    "reuses_state",
     "rmsprop",
     "rmsprop_updates",
     "rprop",
diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py
index 4f9c3dd..9fc7940 100644
--- a/pytensor_ml/optim/alias.py
+++ b/pytensor_ml/optim/alias.py
@@ -6,7 +6,6 @@
     Parameter,
     Transform,
     Updates,
-    reuses_state,
 )
 from pytensor_ml.optim.rules import (
     _require_numeric_learning_rate,
@@ -70,15 +69,10 @@ def sgd(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    # Built once here rather than per invocation, so the velocity it owns is the same buffer on every
-    # step compiled from this rule instead of a fresh one each time.
-    momentum_trace = trace(momentum, nesterov) if momentum else None
-
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        if momentum_trace is None:
+        if not momentum:
             return sgd_updates(
                 loss_gradients_or_updates,
                 parameters,
@@ -91,7 +85,7 @@ def rule(
             learning_rate=1.0,
             namespace=namespace,
         )
-        updates = momentum_trace(updates, parameters)
+        updates = trace(momentum, nesterov)(updates, parameters)
         return scale(learning_rate, namespace=namespace)(updates, parameters)
 
     return rule
@@ -132,7 +126,6 @@ def adam(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -188,7 +181,6 @@ def adamw(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -243,7 +235,6 @@ def nadam(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -295,7 +286,6 @@ def adamax(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -350,7 +340,6 @@ def rprop(
     """
     _require_numeric_learning_rate(learning_rate)
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -403,7 +392,6 @@ def rmsprop(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -453,7 +441,6 @@ def adagrad(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -501,7 +488,6 @@ def adadelta(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 62b840b..a6deb80 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -1,6 +1,4 @@
 from collections.abc import Callable, Sequence
-from contextvars import ContextVar
-from functools import wraps
 
 import numpy as np
 import pytensor
@@ -96,10 +94,8 @@ class Steps(Updates):
 ``parameters`` -- rewriting the rest would halve a clock's advance as readily as a step.
 
 One that keeps state of its own allocates it with :func:`state_for` and takes a ``namespace``, so that
-two of them in one chain write to separate buffers rather than colliding at the serialization boundary.
-Wrap it in :func:`reuses_state` to hold those buffers across invocations when it is used outside a chain;
-:func:`chain` gives each of its own members a frame already.
-
+two of them in one chain are told apart at the serialization boundary. Every invocation allocates afresh.
+The updates dict it returns holds the state, so two functions that share state are compiled from one dict.
 
 .. code-block:: python
 
@@ -467,94 +463,6 @@ def _unreachable_parameter_names(
     ]
 
 
-# A per-parameter slot, or the bare name of a rule-wide counter.
-type _StateKey = tuple[Parameter, str] | str
-
-# Bound by reuses_state for the duration of one rule invocation; None means "allocate fresh".
-_state_buffers: ContextVar[dict[_StateKey, Parameter] | None] = ContextVar(
-    "optimizer_state_buffers", default=None
-)
-
-# The per-parameter slots claimed so far in the current invocation. Reuse *across* invocations is the
-# whole point of the buffers, but two claims on one slot within a single invocation are two components
-# allocating over each other, which is otherwise invisible.
-_claimed_slots: ContextVar[set[_StateKey] | None] = ContextVar(
-    "optimizer_claimed_slots", default=None
-)
-
-
-def reuses_state[**P, R](builds_updates: Callable[P, R]) -> Callable[P, R]:
-    """
-    Give ``builds_updates`` a private set of optimizer-state buffers, reused on every invocation.
-
-    A configured rule such as ``adam(1e-3)`` reads as a value, so it is natural to compile two training
-    functions from one. Without this, each invocation allocates fresh momentum under the *same* derived
-    name: the two steps then share parameters but not optimizer state, which is silently wrong at runtime
-    and raises only later when both are checkpointed together. The same holds for a composed transform,
-    whose state is likewise allocated per call.
-
-    The buffers are keyed per wrapped callable rather than globally so two independently configured
-    optimizers stay independent. They are bound dynamically because :func:`state_for` is reached several
-    call layers below, and threading a cache down would touch every one of them. Nesting is safe: an inner
-    scope restores the outer one on exit, so a rule's own buffers and its enclosing chain's coexist.
-
-    Parameters
-    ----------
-    builds_updates : callable
-        A rule or transform to wrap. Its buffers live as long as the wrapper does.
-
-    Returns
-    -------
-    with_persistent_state : callable
-        ``builds_updates`` with a buffer scope of its own, matching its signature.
-
-    Examples
-    --------
-    Wrap a hand-written transform that allocates state, so two functions compiled from it drive the same
-    buffers rather than each getting fresh ones. :func:`chain` already does this for its own members:
-
-    .. code-block:: python
-
-        from pytensor_ml.optim import reuses_state, state_for, to_updates
-
-
-        def smooth(decay, namespace="smooth"):
-            @reuses_state
-            def transform(loss_gradients_or_updates, parameters):
-                updates = to_updates(loss_gradients_or_updates, parameters)
-                smoothed = updates.copy()
-                for parameter in parameters:
-                    velocity = state_for(parameter, f"{namespace}/velocity")
-                    smoothed[velocity] = decay * velocity + (updates[parameter] - parameter)
-                    smoothed[parameter] = parameter + smoothed[velocity]
-                return smoothed
-
-            return transform
-    """
-    buffers: dict[_StateKey, Parameter] = {}
-
-    @wraps(builds_updates)
-    def with_persistent_state(*args: P.args, **kwargs: P.kwargs) -> R:
-        token = _state_buffers.set(buffers)
-        claimed_token = _claimed_slots.set(set())
-        try:
-            return builds_updates(*args, **kwargs)
-        finally:
-            _state_buffers.reset(token)
-            _claimed_slots.reset(claimed_token)
-
-    return with_persistent_state
-
-
-def _reuse_or_allocate(key: _StateKey, allocate: Callable[[], Parameter]) -> Parameter:
-    buffers = _state_buffers.get()
-    if buffers is None:
-        return allocate()
-    if key not in buffers:
-        buffers[key] = allocate()
-    return buffers[key]
-
-
 def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Parameter:
     """
     Return the optimizer-state shared variable shaped and typed like ``parameter``.
@@ -564,9 +472,8 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param
     callers hold the returned object directly, and reuse within a rule is keyed on the parameter object, so
     two same-named parameters still get distinct buffers rather than silently sharing one.
 
-    Allocates unless the enclosing rule was wrapped in :func:`reuses_state` and already holds this slot.
-    Within one invocation a slot belongs to one component: a second claim on it raises rather than handing
-    two components the same buffer, which only the later writer's updates would survive.
+    A fresh variable on every call. The updates dict a rule returns holds it, so two training functions
+    that share this state are compiled from one updates dict.
 
     Parameters
     ----------
@@ -580,7 +487,7 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param
     Returns
     -------
     state : shared tensor variable
-        The buffer for this slot, allocated on the first claim and returned again on later ones.
+        A new buffer for this slot.
 
     Examples
     --------
@@ -603,46 +510,30 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param
             "parameter names to identify their state at serialization boundaries; give the parameter a name."
         )
 
-    key = (parameter, slot)
-    claimed = _claimed_slots.get()
-    if claimed is not None:
-        if key in claimed:
-            raise ValueError(
-                f"Two components asked for the {slot!r} state of {parameter.name!r} in one step, so the "
-                "second would allocate over the first and only its writes would survive. Give one of them "
-                "a `namespace` of its own, or wrap it in `reuses_state` so it keeps its own buffers."
-            )
-        claimed.add(key)
-
-    def allocate() -> Parameter:
-        value = parameter.get_value(borrow=True)
-        state = pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}")
-        # Keeps `Linear_1_W` and `Linear_1_W/adam/first_moment` numbered onto the same layer.
-        state.layer_name = getattr(parameter, "layer_name", None)
-        return state
-
-    return _reuse_or_allocate(key, allocate)
+    value = parameter.get_value(borrow=True)
+    state = pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}")
+    # Keeps `Linear_1_W` and `Linear_1_W/adam/first_moment` numbered onto the same layer.
+    state.layer_name = getattr(parameter, "layer_name", None)
+    return state
 
 
 def counter(name: str) -> Parameter:
     """
-    Return the training clock a component counts its own steps on.
+    Allocate the training clock a component counts its own steps on.
 
-    Reused across invocations of a rule wrapped in :func:`reuses_state`, so the count keeps advancing. A
-    :class:`~pytensor_ml.params.StepCounter` rather than a plain shared variable, so a schedule can read
+    A :class:`~pytensor_ml.params.StepCounter` rather than a plain shared variable, so a schedule can read
     the same notion of time the rule uses, and :func:`~pytensor_ml.pytensorf.collect_clock_updates` advances
     it for a caller who does not write the advance themselves.
 
     Parameters
     ----------
     name : str
-        Name of the clock, used to match it at serialization boundaries. Two components given the same
-        name share one clock, which is how a rule and the schedule driving it count the same steps.
+        Name of the clock, used to match it at serialization boundaries.
 
     Returns
     -------
     clock : StepCounter
-        The step counter under ``name``, allocated on first use and returned again after that.
+        A new step counter under ``name``.
 
     Examples
     --------
@@ -655,13 +546,12 @@ def counter(name: str) -> Parameter:
 
         rate = cosine_schedule(3e-4, total_steps=10_000)(counter("my_transform/step_count"))
     """
-    return _reuse_or_allocate(name, lambda: step_counter(name))
+    return step_counter(name)
 
 
 def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
     """
-    Return a floatX scalar shared variable, reused across invocations of a rule wrapped in
-    :func:`reuses_state`.
+    Allocate a floatX scalar shared variable.
 
     Parameters
     ----------
@@ -681,10 +571,7 @@ def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
 
         scale = scalar_state("plateau/scale", fill_value=1.0)
     """
-    return _reuse_or_allocate(
-        name,
-        lambda: pytensor.shared(np.asarray(fill_value, dtype=pytensor.config.floatX), name=name),
-    )
+    return pytensor.shared(np.asarray(fill_value, dtype=pytensor.config.floatX), name=name)
 
 
 def require_unique_state_names(updates: Updates) -> None:
@@ -731,9 +618,6 @@ def chain(*transforms: Transform) -> Transform:
 
     A chain is itself a transform, so one composes into another and the result is flat.
 
-    The composed callable owns one set of optimizer-state buffers however many times it is invoked, so two
-    training functions compiled from one chain share its momentum rather than each allocating their own.
-
     Parameters
     ----------
     *transforms : Transform
@@ -775,17 +659,11 @@ def chain(*transforms: Transform) -> Transform:
     if not transforms:
         raise ValueError("chain needs at least one transform.")
 
-    # Each member gets a buffer frame of its own, made once here. Without it a transform that allocates
-    # state without wrapping itself falls through to the chain's frame, where a second such transform
-    # would claim the same slot and quietly take it over.
-    staged = tuple(reuses_state(transform) for transform in transforms)
-
-    @reuses_state
     def combined(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
-        updates = staged[0](loss_gradients_or_updates, parameters)
-        for transform in staged[1:]:
+        updates = transforms[0](loss_gradients_or_updates, parameters)
+        for transform in transforms[1:]:
             updates = transform(updates, parameters)
         return updates
 
diff --git a/pytensor_ml/optim/guards.py b/pytensor_ml/optim/guards.py
index 340ca44..e0740ba 100644
--- a/pytensor_ml/optim/guards.py
+++ b/pytensor_ml/optim/guards.py
@@ -11,7 +11,6 @@
     Parameter,
     Transform,
     Updates,
-    reuses_state,
     scalar_state,
 )
 from pytensor_ml.optim.checks import checked_scalar
@@ -289,7 +288,6 @@ def skip_if(
     elif not isinstance(condition, SkipCondition):
         condition = SkipCondition(condition)
 
-    @reuses_state
     def guarded(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
diff --git a/pytensor_ml/optim/policy.py b/pytensor_ml/optim/policy.py
index 411fd8a..a725f4f 100644
--- a/pytensor_ml/optim/policy.py
+++ b/pytensor_ml/optim/policy.py
@@ -12,7 +12,6 @@
     Steps,
     Transform,
     Updates,
-    reuses_state,
     scalar_state,
 )
 from pytensor_ml.optim.checks import checked_scalar
@@ -134,7 +133,6 @@ def reduce_on_plateau(
         condition_fn=lambda value: value >= 1,
     )
 
-    @reuses_state
     def policy(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
diff --git a/pytensor_ml/optim/train.py b/pytensor_ml/optim/train.py
index 4c86cc2..6dc5837 100644
--- a/pytensor_ml/optim/train.py
+++ b/pytensor_ml/optim/train.py
@@ -36,7 +36,7 @@ def _inconsistent_update(updates: Updates, variable: SharedVariable, new_value:
 
 def compile_train(
     loss: TensorVariable,
-    rule: Transform,
+    rule: Transform | Updates,
     *,
     parameters: Sequence[Parameter] | None = None,
     inputs: Sequence[Variable] | None = None,
@@ -56,8 +56,11 @@ def compile_train(
     ----------
     loss : TensorVariable
         Scalar loss to minimize.
-    rule : Transform
-        A configured optimizer ``(loss_gradients_or_updates, parameters) -> Updates``, e.g. ``adam(1e-3)``.
+    rule : Transform or Updates
+        A configured optimizer ``(loss_gradients_or_updates, parameters) -> Updates``, e.g. ``adam(1e-3)``,
+        or the updates dict one invocation of it returned. A rule allocates its state afresh on every
+        invocation, so two training functions that should share momentum are compiled from one dict:
+        ``updates = adam(1e-3)(loss, parameters)``, then ``compile_train(loss, updates, ...)`` twice.
     parameters : sequence of shared tensor variable, optional
         Parameters to optimize. Collected from ``loss`` with :func:`collect_differentiable_params` when
         omitted, so parameters the loss detaches with a stop-gradient are left alone. A detached parameter
@@ -132,7 +135,7 @@ def compile_train(
     if inputs is None:
         inputs = collect_data_inputs([loss, *extra_outputs, *extra_updates.values()])
 
-    result = rule(loss, parameters)
+    result = rule(loss, parameters) if callable(rule) else rule
     if isinstance(result, Gradients):
         raise ValueError(
             "The rule returned gradients rather than the steps to take, so every parameter would move "
diff --git a/pytensor_ml/optim/transform.py b/pytensor_ml/optim/transform.py
index 27aebd7..71b997b 100644
--- a/pytensor_ml/optim/transform.py
+++ b/pytensor_ml/optim/transform.py
@@ -12,7 +12,6 @@
     Transform,
     Updates,
     read_rate,
-    reuses_state,
     state_for,
     steps_of,
     to_updates,
@@ -76,7 +75,6 @@ def trace(decay: float = 0.9, nesterov: bool = False, *, namespace: str = "trace
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def transform(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -137,7 +135,6 @@ def scale(factor: LearningRate, *, namespace: str = "scale") -> Transform:
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
 
-    @reuses_state
     def transform(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
diff --git a/tests/optim/test_composition.py b/tests/optim/test_composition.py
index ff14c46..02c68fc 100644
--- a/tests/optim/test_composition.py
+++ b/tests/optim/test_composition.py
@@ -1,5 +1,6 @@
 from itertools import pairwise
 
+import cloudpickle
 import numpy as np
 import pytensor.tensor as pt
 import pytest
@@ -189,3 +190,18 @@ def test_two_scheduled_scales_sharing_a_namespace_collide_loudly():
         scale_by_schedule(schedule, namespace="decay"),
     )
     compile_train(loss, apart)
+
+
+def test_an_unused_rule_ships_through_cloudpickle_and_trains_after():
+    """A fit config that carries a configured rule to a remote worker goes through cloudpickle, which
+    serializes a nested function's globals by value. A rule holds hyperparameters and nothing else, so the
+    whole spine ships, and the state it allocates on the worker is the worker's."""
+    rule = skip_if(chain(adam(cosine_schedule(1e-2, 1000)), clip_by_global_norm(10.0)))
+    restored = cloudpickle.loads(cloudpickle.dumps(rule))
+
+    p, loss = quadratic_problem()
+    step = compile_train(loss, restored)
+    before = step(GOOD)
+    step(GOOD)
+
+    assert step(GOOD) < before
diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index 06295af..9dcbb9d 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -318,12 +318,12 @@ def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
     only one of them leaves the two measuring different times."""
     p, loss = quadratic_problem()
     their_clock = step_counter("their_clock")
-    rule = adam(learning_rate=0.1)
-    adam_clock = next(key for key in rule(loss, [p]) if key.name == "adam/step_count")
+    updates = adam(learning_rate=0.1)(loss, [p])
+    adam_clock = next(key for key in updates if key.name == "adam/step_count")
     adam_clock.set_value(np.asarray(37, dtype="int64"))
 
     with pytest.raises(ValueError, match="hold different step counts"):
-        compile_train(loss, rule, extra_outputs=[their_clock.astype("float64")])
+        compile_train(loss, updates, extra_outputs=[their_clock.astype("float64")])
 
 
 @pytest.mark.parametrize(
diff --git a/tests/optim/test_policy.py b/tests/optim/test_policy.py
index 5c146c9..121a0a7 100644
--- a/tests/optim/test_policy.py
+++ b/tests/optim/test_policy.py
@@ -214,9 +214,9 @@ def test_the_decision_is_made_on_the_window_mean_not_its_last_loss():
     the loss is exactly what is fed and the policy's own state is all that changes."""
     loss = pt.scalar("loss")
     scale = scalar_state("plateau/scale", fill_value=1.0)
-    rule = reduce_on_plateau(adam(learning_rate=scale), scale, accumulation_size=2)
-    best_loss = next(v for v in rule(loss, []) if v.name == "plateau/best_loss")
-    step = compile_train(loss, rule, parameters=[], inputs=[loss])
+    updates = reduce_on_plateau(adam(learning_rate=scale), scale, accumulation_size=2)(loss, [])
+    best_loss = next(v for v in updates if v.name == "plateau/best_loss")
+    step = compile_train(loss, updates, parameters=[], inputs=[loss])
 
     step(np.asarray(4.0, dtype=config.floatX))
     step(np.asarray(0.0, dtype=config.floatX))
diff --git a/tests/optim/test_public_surface.py b/tests/optim/test_public_surface.py
index 6a1fbf6..1fcda86 100644
--- a/tests/optim/test_public_surface.py
+++ b/tests/optim/test_public_surface.py
@@ -16,9 +16,7 @@ def test_exported_names_are_unique_and_sorted():
     assert optim.__all__ == sorted(set(optim.__all__))
 
 
-@pytest.mark.parametrize(
-    "name", ["to_updates", "steps_of", "state_for", "reuses_state", "counter", "chain"]
-)
+@pytest.mark.parametrize("name", ["to_updates", "steps_of", "state_for", "counter", "chain"])
 def test_the_transform_authoring_api_stays_public(name):
     """The ``Transform`` docstring tells authors to write a transform with these. Dropping one from
     ``__all__`` would leave the documented path running through ``pytensor_ml.optim.base``."""
diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py
index 9ed0aa7..8151575 100644
--- a/tests/optim/test_rules.py
+++ b/tests/optim/test_rules.py
@@ -295,11 +295,10 @@ def test_sgd_names_the_step_counter_its_schedule_reads():
     [lambda: adam(learning_rate=1e-2), lambda: sgd(learning_rate=1e-2, momentum=0.9)],
     ids=["state_from_a_rule", "state_from_a_transform"],
 )
-def test_reused_rule_shares_its_optimizer_state(make_rule):
-    """A configured rule reads as a value, so compiling two training functions from one is natural. Both
-    must drive the same buffers: separate ones under the same derived name are silently wrong at runtime,
-    and collide only later when both are checkpointed together. Momentum SGD is included because its
-    velocity comes from a transform rather than the rule, which is a separate allocation path."""
+def test_each_invocation_of_a_rule_allocates_its_own_state(make_rule):
+    """A rule builds a graph, and the graph holds the state. Two invocations are two optimizers, so the
+    buffers they allocate are distinct objects even though their derived names agree. Momentum SGD is
+    included because its velocity comes from a transform rather than the rule, a separate allocation path."""
     p = trainable(np.zeros(3), name="w")
     loss = (p**2).sum()
     rule = make_rule()
@@ -307,21 +306,22 @@ def test_reused_rule_shares_its_optimizer_state(make_rule):
     first = {key for key in rule(loss, [p]) if key is not p}
     second = {key for key in rule(loss, [p]) if key is not p}
 
-    assert first and first == second
+    assert first and not first & second
 
 
-def test_two_functions_from_one_rule_continue_the_same_momentum():
-    """What the shared buffers buy: the second function continues the first's trajectory instead of
-    restarting it. Under a constant gradient, momentum SGD's step at iteration ``t`` is
-    ``lr * g * (1 - m**t) / (1 - m)``, so a continued second step is 1.9x a restarted one at ``m = 0.9``."""
+def test_two_functions_from_one_updates_dict_continue_the_same_momentum():
+    """Two training functions share state by being compiled from one updates dict. The second then
+    continues the first's trajectory instead of restarting it. Under a constant gradient, momentum SGD's
+    step at iteration ``t`` is ``lr * g * (1 - m**t) / (1 - m)``, so a continued second step is 1.9x a
+    restarted one at ``m = 0.9``."""
     p = trainable(np.zeros(2), name="w")
     gradient = np.array([2.0, -0.5])
     loss = (pt.constant(gradient, dtype=floatX) * p).sum()  # constant gradient, independent of p
     learning_rate, momentum = 0.1, 0.9
-    rule = sgd(learning_rate=learning_rate, momentum=momentum)
+    updates = sgd(learning_rate=learning_rate, momentum=momentum)(loss, [p])
 
-    step_once = function([], loss, updates=rule(loss, [p]))
-    step_again = function([], loss, updates=rule(loss, [p]))
+    step_once = function([], loss, updates=updates)
+    step_again = function([], loss, updates=updates)
 
     step_once()
     before = p.get_value().copy()
@@ -331,18 +331,6 @@ def test_two_functions_from_one_rule_continue_the_same_momentum():
     np.testing.assert_allclose(p.get_value() - before, continued, rtol=RTOL)
 
 
-def test_separately_configured_rules_keep_independent_state():
-    """Buffers are memoized per rule, not globally, so two optimizers over the same parameter do not
-    quietly train through each other's momentum."""
-    p = trainable(np.zeros(3), name="w")
-    loss = (p**2).sum()
-
-    first = {key for key in adam(learning_rate=1e-2)(loss, [p]) if key is not p}
-    second = {key for key in adam(learning_rate=1e-2)(loss, [p]) if key is not p}
-
-    assert not first & second
-
-
 def test_adamw_first_step_applies_decoupled_decay():
     """AdamW adds a decoupled decay term to Adam's sign-descent step: the first-step displacement is
     ``-lr * (sign(g) + weight_decay * p)``. At t = 1 bias correction makes the Adam part ``sign(g)`` (the
diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py
index 3171655..3409640 100644
--- a/tests/optim/test_training.py
+++ b/tests/optim/test_training.py
@@ -375,16 +375,15 @@ def test_extra_updates_write_state_no_gradient_produces():
 
 
 def test_extra_updates_accept_a_write_that_agrees_with_the_rule():
-    """Two components writing one variable is only a problem when they disagree. A rule builds a fresh
-    expression on every invocation, so an extra update carrying a structurally identical one is the same
-    write said twice -- which is how components share a quantity rather than fight over it."""
+    """Two components writing one variable is only a problem when they disagree. An extra update carrying
+    a structurally identical expression is the same write said twice -- which is how components share a
+    quantity rather than fight over it."""
     p = trainable(np.array([2.0]), name="w")
     loss = 0.5 * (p**2).sum()
-    rule = adam(1e-1)
-    first_invocation = rule(loss, [p])
-    moment = next(key for key in first_invocation if key.name == "w/adam/first_moment")
+    updates = adam(1e-1)(loss, [p])
+    moment = next(key for key in updates if key.name == "w/adam/first_moment")
 
-    step = compile_train(loss, rule, extra_updates={moment: first_invocation[moment]}, inputs=[])
+    step = compile_train(loss, updates, extra_updates={moment: updates[moment]}, inputs=[])
     step()
 
     assert not np.allclose(
@@ -418,11 +417,11 @@ def test_extra_updates_reject_a_write_the_rule_already_makes():
     # 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")
+    updates = adam(1e-1)(loss, [p])
+    first_moment = next(key for key in updates 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=[])
+        compile_train(loss, updates, extra_updates={first_moment: first_moment * 0.0}, inputs=[])
 
 
 def test_extra_updates_reject_a_write_the_model_already_makes():
diff --git a/tests/optim/test_transform.py b/tests/optim/test_transform.py
index 746b88b..020769b 100644
--- a/tests/optim/test_transform.py
+++ b/tests/optim/test_transform.py
@@ -54,20 +54,6 @@ def test_scale_by_schedule_allocates_a_clock_per_namespace():
     ]
 
 
-def test_scale_by_schedule_reuses_its_clock_across_invocations():
-    """Two functions compiled from one configured transform must read one clock; a second would restart the
-    schedule at zero while the first kept counting."""
-    p = trainable(np.zeros(1), name="w")
-    transform = scale_by_schedule(linear_schedule(1.0, total_steps=4))
-
-    (first_clock,) = collect_step_counters(transform({p: p + pt.constant(np.array([1.0]))}, [p])[p])
-    (second_clock,) = collect_step_counters(
-        transform({p: p + pt.constant(np.array([1.0]))}, [p])[p]
-    )
-
-    assert first_clock is second_clock
-
-
 def test_trace_accumulates_velocity_with_decay():
     p = trainable(np.zeros(1), name="w")
     out = trace(0.9)({p: p + pt.constant(np.array([1.0]))}, [p])
@@ -116,19 +102,6 @@ def test_a_rule_headed_chain_composes_again_as_a_rule():
     np.testing.assert_allclose(function([], out[p])(), [1.5])
 
 
-def test_chain_reuses_transform_state_across_invocations():
-    """A chain's transforms allocate state too, so two functions compiled from one chain must share it.
-    Without reuse the second function silently restarts with its own velocity buffer."""
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()
-    rule = chain(trace(0.9), scale(0.1))
-
-    first = {key for key in rule(sgd_updates(loss, [p], learning_rate=1.0), [p]) if key is not p}
-    second = {key for key in rule(sgd_updates(loss, [p], learning_rate=1.0), [p]) if key is not p}
-
-    assert first and first == second
-
-
 def test_separately_configured_chains_keep_independent_state():
     p = trainable(np.array([2.0]), name="w")
     loss = 0.5 * (p**2).sum()
diff --git a/tests/optim/test_update_space.py b/tests/optim/test_update_space.py
index 3bdd582..792cef2 100644
--- a/tests/optim/test_update_space.py
+++ b/tests/optim/test_update_space.py
@@ -62,8 +62,7 @@ def run_after_spike(rule, n_steps=STEPS_AFTER_SPIKE + 1, spike=1e6):
 
 
 def stateful_transform(decay, slot="smooth/velocity"):
-    """A hand-written stateful transform, written the way the ``Transform`` docstring teaches: with
-    ``state_for`` and without ``reuses_state``, since nothing at the call site suggests the decorator."""
+    """A hand-written stateful transform, written the way the ``Transform`` docstring teaches."""
 
     def transform(loss_gradients_or_updates, parameters):
         updates = to_updates(loss_gradients_or_updates, parameters)
@@ -346,28 +345,8 @@ def test_a_rule_behind_another_rule_is_refused(build_rule, named):
         compile_train(loss, build_rule(), parameters=[parameter])
 
 
-@pytest.mark.parametrize(
-    "build_rule",
-    [lambda: adam(1e-3), lambda: trace(0.9), lambda: sgd(0.1, momentum=0.9)],
-    ids=["adam", "trace", "momentum-sgd"],
-)
-def test_one_configured_transform_keeps_one_set_of_buffers(build_rule):
-    """Compiling two training functions from one configured transform is natural, and both have to drive
-    the same state: separate buffers under one derived name are wrong at runtime and collide only later,
-    when both are checkpointed together."""
-    parameter = trainable(np.zeros(3, dtype=config.floatX), name="w")
-    gradient = pt.constant(np.ones(3, dtype=config.floatX))
-    rule = build_rule()
-
-    first = {key for key in rule([gradient], [parameter]) if key is not parameter}
-    second = {key for key in rule([gradient], [parameter]) if key is not parameter}
-
-    assert first and first == second
-
-
 def test_a_hand_written_stateful_transform_owns_its_buffers_in_a_chain():
-    """A chain gives each member a frame of its own. Without one, a transform that allocates state falls
-    through to the chain's frame, where the next such transform takes the slot over."""
+    """Two hand-written transforms in one chain each allocate their own buffers under their own names."""
     parameter = trainable(np.zeros(3, dtype=config.floatX), name="w")
     loss = (parameter**2).sum()
 
@@ -383,35 +362,6 @@ def test_a_hand_written_stateful_transform_owns_its_buffers_in_a_chain():
     ]
 
 
-def test_a_hand_written_stateful_transform_keeps_its_buffers_across_invocations():
-    """Two training functions compiled from one chain must drive the same state, which is the whole reason
-    the buffers are held rather than reallocated."""
-    parameter = trainable(np.zeros(3, dtype=config.floatX), name="w")
-    loss = (parameter**2).sum()
-    rule = chain(stateful_transform(0.9))
-
-    first = {id(key) for key in rule(loss, [parameter]) if key.name and "velocity" in key.name}
-    again = {id(key) for key in rule(loss, [parameter]) if key.name and "velocity" in key.name}
-
-    assert first and first == again
-
-
-def test_two_claims_on_one_slot_in_a_single_step_are_refused():
-    """Reuse across invocations is the point of the buffers; two claims within one is two components
-    allocating over each other, and only the second one's writes would survive."""
-    parameter = trainable(np.zeros(3, dtype=config.floatX), name="w")
-
-    def claims_twice(loss_gradients_or_updates, parameters):
-        updates = to_updates(loss_gradients_or_updates, parameters)
-        for each in parameters:
-            state_for(each, "shared/velocity")
-            state_for(each, "shared/velocity")
-        return updates
-
-    with pytest.raises(ValueError, match="in one step"):
-        chain(claims_twice, sgd(0.1))((parameter**2).sum(), [parameter])
-
-
 def test_two_transforms_of_one_kind_need_distinct_namespaces():
     """Same-kind transforms allocate under the same derived name, so they collide at the serialization
     boundary until one is given a namespace -- and the error has to say that rather than blame parameters."""

From cb9af7094034549f11f3eb87ff5e278c10b263e1 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:28:18 +0800
Subject: [PATCH 03/16] Keep only the gradients a rule descends

---
 pytensor_ml/optim/base.py    | 17 +++++++++++++++--
 tests/optim/test_training.py | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index a6deb80..6ac39e0 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -10,7 +10,7 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.sharedvar import TensorSharedVariable
 
-from pytensor_ml.params import step_counter
+from pytensor_ml.params import TrainableParameter, step_counter
 from pytensor_ml.pytensorf import rewrite_pregrad
 
 type Parameter = TensorSharedVariable
@@ -426,7 +426,20 @@ def gradients_to_descend(
             "descends along what it reads, so it would negate that step and move the parameters uphill. "
             "Keep one rule in a chain and shape its step with `scale`, `trace`, or a clip after it."
         )
-    return incoming, steps_of(incoming, parameters)
+    gradients = steps_of(incoming, parameters)
+    if isinstance(incoming, Gradients):
+        # A gradient for a parameter this rule does not descend cannot travel on inside a Steps dict: it
+        # would compile as `p + g`, an ascent. Another rule over that parameter reads it from the same
+        # Gradients dict this one did, so dropping it here loses nothing.
+        own_parameters = set(parameters)
+        incoming = Gradients(
+            {
+                key: value
+                for key, value in incoming.items()
+                if key in own_parameters or not isinstance(key, TrainableParameter)
+            }
+        )
+    return incoming, gradients
 
 
 def steps_of(updates: Updates, parameters: Sequence[Parameter]) -> list[TensorVariable]:
diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py
index 3409640..a90cdfe 100644
--- a/tests/optim/test_training.py
+++ b/tests/optim/test_training.py
@@ -10,7 +10,15 @@
 from pytensor_ml.activations import LeakyReLU, Tanh
 from pytensor_ml.layers import BatchNorm, Linear, Sequential
 from pytensor_ml.loss import CrossEntropy, SquaredError, supervised_loss
-from pytensor_ml.optim import adam, adamw, compile_train, cosine_schedule, sgd
+from pytensor_ml.optim import (
+    adam,
+    adamw,
+    chain,
+    clip_by_global_norm,
+    compile_train,
+    cosine_schedule,
+    sgd,
+)
 from pytensor_ml.optim.base import state_for
 from pytensor_ml.params import step_counter, trainable
 from pytensor_ml.pytensorf import collect_non_trainable_params, collect_trainable_params
@@ -242,6 +250,29 @@ def decay_the_weights_only(loss_or_gradients, parameters):
     np.testing.assert_allclose(bias.get_value(), [3.0 - 0.1], rtol=1e-3)
 
 
+def test_a_rule_keeps_only_the_gradients_it_descends():
+    """Behind a gradient-space transform, each per-group rule is handed every parameter's gradient. A rule
+    that passed the others on inside its Steps dict would have them compile as ``p + g``, an ascent, and
+    the merge below would keep whichever copy came last."""
+    weight = trainable(np.array([2.0]), name="weight")
+    bias = trainable(np.array([3.0]), name="bias")
+    loss = 0.5 * ((weight**2).sum() + (bias**2).sum())
+
+    def per_group(loss_or_gradients, parameters):
+        return {
+            **adam(learning_rate=0.1)(loss_or_gradients, [bias]),
+            **adamw(learning_rate=0.1)(loss_or_gradients, [weight]),
+        }
+
+    step = compile_train(
+        loss, chain(clip_by_global_norm(100.0), per_group), parameters=[weight, bias], inputs=[]
+    )
+    step()
+
+    assert weight.get_value() < 2.0
+    assert bias.get_value() < 3.0
+
+
 def test_compile_train_includes_non_trainable_updates():
     # compile_train merges batch-norm running-stat updates that a bare gradient rule would omit.
     X = pt.tensor("X", shape=(None, 4))

From 55e6651c008cd59983d883bbafb2cfa07914982e Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:28:37 +0800
Subject: [PATCH 04/16] Refuse an updates dict that leaves a parameter
 unwritten

---
 pytensor_ml/optim/train.py   |  8 ++++++++
 tests/optim/test_training.py | 10 ++++++++++
 2 files changed, 18 insertions(+)

diff --git a/pytensor_ml/optim/train.py b/pytensor_ml/optim/train.py
index 6dc5837..40b9323 100644
--- a/pytensor_ml/optim/train.py
+++ b/pytensor_ml/optim/train.py
@@ -185,6 +185,14 @@ def compile_train(
                 "differently."
             )
 
+    unwritten = [parameter.name for parameter in parameters if parameter not in updates]
+    if unwritten:
+        raise ValueError(
+            f"No update reaches {unwritten}. Every parameter collected from the loss, or given in "
+            "`parameters`, needs an entry in the updates, so an updates dict built over a subset of them "
+            "would silently freeze the rest. Leave a parameter out of `parameters` to freeze it on purpose."
+        )
+
     require_unique_state_names(updates)
 
     outputs = [loss, *extra_outputs] if extra_outputs else loss
diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py
index a90cdfe..b4d9b32 100644
--- a/tests/optim/test_training.py
+++ b/tests/optim/test_training.py
@@ -273,6 +273,16 @@ def per_group(loss_or_gradients, parameters):
     assert bias.get_value() < 3.0
 
 
+def test_an_updates_dict_missing_a_collected_parameter_is_refused():
+    """A dict built over a subset of the parameters would leave the rest untouched with nothing to say so."""
+    weight = trainable(np.array([2.0]), name="weight")
+    bias = trainable(np.array([3.0]), name="bias")
+    loss = 0.5 * ((weight**2).sum() + (bias**2).sum())
+
+    with pytest.raises(ValueError, match=r"No update reaches \['bias'\]"):
+        compile_train(loss, adam(learning_rate=0.1)(loss, [weight]), inputs=[])
+
+
 def test_compile_train_includes_non_trainable_updates():
     # compile_train merges batch-norm running-stat updates that a bare gradient rule would omit.
     X = pt.tensor("X", shape=(None, 4))

From b0faa711ff572e01999eb1edeec194340b35cdd5 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:29:29 +0800
Subject: [PATCH 05/16] Check optimizer state for finiteness, not only the
 parameters

---
 pytensor_ml/optim/guards.py | 14 +++++++++-----
 tests/optim/test_guards.py  | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/pytensor_ml/optim/guards.py b/pytensor_ml/optim/guards.py
index e0740ba..c8ed25e 100644
--- a/pytensor_ml/optim/guards.py
+++ b/pytensor_ml/optim/guards.py
@@ -1,6 +1,7 @@
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 
+import numpy as np
 import pytensor.tensor as pt
 
 from pytensor.raise_op import CheckAndRaise
@@ -94,16 +95,18 @@ def __call__(self, updates: Updates, parameters: Sequence[Parameter]) -> TensorV
 
 def nonfinite() -> SkipCondition:
     """
-    Throw the step away when any parameter the rule would write is inf or NaN.
+    Throw the step away when any value the rule would write is inf or NaN.
 
     The condition behind :func:`apply_if_finite`, and the one to reach for when there is no scale to
     threshold on. It fires only once a value has already gone non-finite, which on a diverging run is a
     lagging alarm: the step before is typically finite and enormous. :func:`large_step` catches that one.
 
-    Only the parameters are checked. A policy is free to keep a sentinel among its own state --
+    Optimizer state is checked along with the parameters, because a poisoned moment need not reach its
+    parameter: an infinite second moment divides the step to zero, which is finite, and the parameter then
+    never moves again. Training clocks are integers and are never checked. A variable that is already
+    non-finite when the guard is built is a sentinel rather than a failure --
     :func:`~pytensor_ml.optim.policy.reduce_on_plateau` holds an infinite best-loss until it has seen a full
-    window -- and that is not a step to throw away. Optimizer state cannot hide a NaN for long in any case,
-    since a poisoned moment reaches its parameter on the very next step.
+    window -- and is left alone.
 
     Examples
     --------
@@ -128,8 +131,9 @@ def nonfinite() -> SkipCondition:
     def decide(updates: Updates, parameters: Sequence[Parameter]) -> TensorVariable:
         checked_values = [
             new_value
-            for new_value in (updates[parameter] for parameter in parameters)
+            for variable, new_value in updates.items()
             if new_value.dtype.startswith("float")
+            and bool(np.all(np.isfinite(variable.get_value(borrow=True))))
         ]
         if not checked_values:
             raise ValueError(
diff --git a/tests/optim/test_guards.py b/tests/optim/test_guards.py
index 99f1a9f..a8b30b2 100644
--- a/tests/optim/test_guards.py
+++ b/tests/optim/test_guards.py
@@ -66,6 +66,25 @@ def test_resumes_training_where_the_skipped_step_left_it():
     np.testing.assert_allclose(p.get_value(), reference_p.get_value())
 
 
+def test_an_infinite_moment_behind_a_finite_parameter_is_skipped():
+    """An overflowed second moment divides adam's step to zero, so the parameter stays finite while the
+    rule is broken for good: every later step is ``finite / inf``. Checking the parameters alone lets that
+    through, and the parameter never moves again."""
+    with config.change_flags(floatX="float32"):
+        gradient = pt.scalar("g", dtype="float32")
+        p = trainable(np.array([2.0], dtype="float32"), name="w")
+        step = compile_train((gradient * p).sum(), apply_if_finite(adam(0.1)))
+
+        step(np.float32(1e30))
+        total_skips = state_named(step, "skip_if/total_skips")
+        second_moment = state_named(step, "w/adam/second_moment")
+
+        assert float(total_skips.get_value()) == 1
+        assert np.isfinite(second_moment.get_value()).all()
+        step(np.float32(1.0))
+        assert p.get_value() < 2.0
+
+
 def test_counts_consecutive_skips_and_resets_on_a_good_step():
     p, loss = poisonable_problem()
     consecutive = scalar_state("skip_if/consecutive_skips")

From 9fdf0f26c8b609829f0c304b7570b3cae5f7291e Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:30:06 +0800
Subject: [PATCH 06/16] Check every counting clock for agreement, exempting
 only pinned ones

---
 pytensor_ml/optim/train.py        | 16 +++-------------
 pytensor_ml/pytensorf/collect.py  | 26 ++++++++++++++------------
 tests/optim/test_learning_rate.py | 13 -------------
 tests/test_collect.py             | 25 ++++++++++++++++++-------
 4 files changed, 35 insertions(+), 45 deletions(-)

diff --git a/pytensor_ml/optim/train.py b/pytensor_ml/optim/train.py
index 40b9323..bf70759 100644
--- a/pytensor_ml/optim/train.py
+++ b/pytensor_ml/optim/train.py
@@ -169,21 +169,11 @@ def compile_train(
 
     # 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.
+    # Assigned per key rather than merged, for the same invariance reason as the statistics above.
     for clock, next_count in collect_clock_updates(
-        [loss, *extra_outputs, *updates.values()]
+        [loss, *extra_outputs, *updates.values()], already_written=updates
     ).items():
-        written = updates.get(clock)
-        if written is None:
-            updates[clock] = next_count
-        elif equal_computations([written], [next_count]):
-            pass  # the rule advances this clock itself, identically, so its own write stands
-        else:
-            raise ValueError(
-                f"The training clock {clock.name!r} is already advanced by an expression that is not the "
-                "one-step advance. A clock advances once per step, so the two writes cannot both take "
-                "effect; drop yours, or write it as `clock + 1` if it is the same advance spelled "
-                "differently."
-            )
+        updates[clock] = next_count
 
     unwritten = [parameter.name for parameter in parameters if parameter not in updates]
     if unwritten:
diff --git a/pytensor_ml/pytensorf/collect.py b/pytensor_ml/pytensorf/collect.py
index ff77c9a..2c71b0a 100644
--- a/pytensor_ml/pytensorf/collect.py
+++ b/pytensor_ml/pytensorf/collect.py
@@ -328,7 +328,7 @@ def collect_step_counters(outputs: Variable | Sequence[Variable]) -> list[StepCo
 
 def collect_clock_updates(
     outputs: Variable | Sequence[Variable],
-    already_written: Container[SharedVariable] = (),
+    already_written: Mapping[SharedVariable, Variable] | None = None,
 ) -> dict[StepCounter, TensorVariable]:
     """
     Collect the advance for every training clock the graph reads, ready to pass as a function's ``updates``.
@@ -340,11 +340,12 @@ def collect_clock_updates(
     ----------
     outputs
         One or more graph outputs to trace back from.
-    already_written : container of shared variable, optional
-        The variables the caller writes themselves, typically the ``updates`` mapping itself. Any clock
-        among them is left out of the result and exempt from the rule that every clock agrees on its step
-        count, since a clock this step does not advance is not counting its steps. Default is empty, which
-        collects and checks every clock the graph reads.
+    already_written : mapping from shared variable to its next value, optional
+        The updates the caller writes themselves, typically the ``updates`` mapping itself. Any clock among
+        them is left out of the result, so a rule's own advance stands. A clock written as itself is pinned:
+        it is not counting this step's steps, so it is also exempt from the rule that every clock agrees
+        on its step count. A clock written any other way is counting, and is checked. Default is empty,
+        which collects and checks every clock the graph reads.
 
     Returns
     -------
@@ -372,18 +373,19 @@ def collect_clock_updates(
 
         clock_updates = collect_clock_updates(list(updates.values()))
     """
-    counters = [
-        counter for counter in collect_step_counters(outputs) if counter not in already_written
-    ]
-    step_counts = {int(counter.get_value()) for counter in counters}
+    already_written = {} if already_written is None else already_written
+    clocks = collect_step_counters(outputs)
+
+    counting = [clock for clock in clocks if already_written.get(clock) is not clock]
+    step_counts = {int(clock.get_value()) for clock in counting}
     if len(step_counts) > 1:
         raise ValueError(
-            f"Training clocks {sorted(str(counter.name) for counter in counters)} hold different step "
+            f"Training clocks {sorted(str(clock.name) for clock in counting)} hold different step "
             f"counts {sorted(step_counts)}. They all count training steps, so this usually means a "
             "checkpoint restored some of them and not the others. Restore all of them, or set them to the "
             "same count before compiling."
         )
-    return {counter: counter.advance() for counter in counters}
+    return {clock: clock.advance() for clock in clocks if clock not in already_written}
 
 
 def collect_non_trainable_updates(
diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index 9dcbb9d..16f3b56 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -232,19 +232,6 @@ def rule(loss_or_gradients, parameters):
     assert int(clock.get_value()) == 2  # advanced once per step, not twice
 
 
-def test_compile_train_rejects_a_rule_that_advances_the_clock_differently():
-    """Two disagreeing writers would leave the clock holding one of them while both looked configured."""
-    _, loss = quadratic_problem()
-    clock = step_counter()
-
-    def rule(loss_or_gradients, parameters):
-        updates = sgd(learning_rate=cosine_schedule(0.1, 10)(clock))(loss_or_gradients, parameters)
-        return {**updates, clock: clock + 5}
-
-    with pytest.raises(ValueError, match="not the one-step advance"):
-        compile_train(loss, rule)
-
-
 def test_a_clock_read_only_by_an_extra_output_still_advances():
     """Collection covers the diagnostics too, or a reported step count would sit at zero."""
     _, loss = quadratic_problem()
diff --git a/tests/test_collect.py b/tests/test_collect.py
index 04b0252..c8bc3c2 100644
--- a/tests/test_collect.py
+++ b/tests/test_collect.py
@@ -221,16 +221,27 @@ def test_clocks_that_disagree_raise(self):
 
     def test_an_already_written_clock_is_left_out_of_the_result(self):
         threaded, pinned = step_counter(name="threaded"), step_counter(name="pinned")
-        assert set(collect_clock_updates([threaded + pinned], already_written=[pinned])) == {
-            threaded
-        }
+        assert set(
+            collect_clock_updates([threaded + pinned], already_written={pinned: pinned})
+        ) == {threaded}
+
+    def test_a_pinned_clock_is_exempt_from_the_agreeing_count_check(self):
+        """A clock written as itself is not counting this step's steps, so it has nothing to agree with."""
+        restored, fresh = step_counter(name="restored"), step_counter(name="fresh")
+        restored.set_value(np.asarray(120, dtype="int64"))
+
+        assert set(
+            collect_clock_updates([restored + fresh], already_written={restored: restored})
+        ) == {fresh}
 
-    def test_an_already_written_clock_is_exempt_from_the_agreeing_count_check(self):
-        """A clock the caller writes is not counting this step's steps, so it has nothing to agree with."""
+    def test_a_clock_the_caller_advances_is_still_checked(self):
+        """A rule writes its own clock's advance. That clock is counting these steps, so a checkpoint that
+        restored it and not the others is the very thing the check exists to catch."""
         restored, fresh = step_counter(name="restored"), step_counter(name="fresh")
         restored.set_value(np.asarray(120, dtype="int64"))
 
-        assert set(collect_clock_updates([restored + fresh], already_written=[restored])) == {fresh}
+        with pytest.raises(ValueError, match="hold different step counts"):
+            collect_clock_updates([restored + fresh], already_written={restored: restored + 1})
 
     def test_pinning_one_clock_does_not_excuse_the_rest_from_agreeing(self):
         """The exemption covers only the clock the caller pinned. Two clocks the step still advances have
@@ -241,7 +252,7 @@ def test_pinning_one_clock_does_not_excuse_the_rest_from_agreeing(self):
         restored.set_value(np.asarray(120, dtype="int64"))
 
         with pytest.raises(ValueError, match="hold different step counts"):
-            collect_clock_updates([pinned + restored + fresh], already_written=[pinned])
+            collect_clock_updates([pinned + restored + fresh], already_written={pinned: pinned})
 
     def test_accepts_a_single_variable_or_a_list(self):
         clock = step_counter()

From c2f9afe75e522b394f3a4034c24def9ee3c32bfe Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:30:40 +0800
Subject: [PATCH 07/16] Hold a rule's clock back through a skipped step

A rule's clock counts the updates it applied, and a skipped step was not one: a count that advanced anyway bias-corrects the next step against a moment that never saw it. A caller-held clock is not the rule's state and keeps counting calls.
---
 pytensor_ml/optim/guards.py     | 18 +++++++-----------
 tests/optim/test_composition.py | 24 ++++++++++++++++++++----
 tests/optim/test_guards.py      | 25 +++++++++----------------
 3 files changed, 36 insertions(+), 31 deletions(-)

diff --git a/pytensor_ml/optim/guards.py b/pytensor_ml/optim/guards.py
index c8ed25e..aab79d4 100644
--- a/pytensor_ml/optim/guards.py
+++ b/pytensor_ml/optim/guards.py
@@ -15,7 +15,6 @@
     scalar_state,
 )
 from pytensor_ml.optim.checks import checked_scalar
-from pytensor_ml.params import StepCounter
 
 type Decision = Callable[[Updates, Sequence[Parameter]], TensorVariable]
 """
@@ -223,8 +222,10 @@ def skip_if(
 
     The guard covers what ``rule`` writes and nothing else: a batch-norm running statistic, written by the
     model and folded in by :func:`~pytensor_ml.optim.train.compile_train` outside the rule, is not held back
-    with the step, and the step still returns the loss that produced the skip. Training clocks are exempt
-    too -- a skipped step still consumed a step, so the schedules reading them advance as usual.
+    with the step, and the step still returns the loss that produced the skip. The rule's own clock is held
+    back with the rest, since it counts the updates the rule applied and a skipped one was not: a count that
+    advanced anyway would bias-correct the next step against a moment that never saw it. A clock the caller
+    holds and passes into a schedule is not the rule's state, so it keeps counting calls, skipped or not.
 
     .. code-block:: python
 
@@ -296,14 +297,9 @@ def guarded(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
         updates = rule(loss_gradients_or_updates, parameters).copy()
-        # Snapshotted before the counters are added, since those are the one thing a skipped step still has
-        # to write: freeze them along with everything else and the guard can never count its way to the
-        # error, which is a silent failure rather than a loud one.
-        held_back = {
-            variable: new_value
-            for variable, new_value in updates.items()
-            if not isinstance(variable, StepCounter)
-        }
+        # Snapshotted before the guard's own counters are added, since those record the skip and are the
+        # one thing a skipped step still writes.
+        held_back = dict(updates)
 
         skipping = condition(updates, parameters)
         consecutive = _counter_or_new(consecutive_skips, f"{namespace}/consecutive_skips")
diff --git a/tests/optim/test_composition.py b/tests/optim/test_composition.py
index 02c68fc..dd9ad7f 100644
--- a/tests/optim/test_composition.py
+++ b/tests/optim/test_composition.py
@@ -109,9 +109,10 @@ def test_a_plateau_policy_composes_with_a_schedule():
     assert int(clock.get_value()) == 6
 
 
-def test_a_guard_keeps_a_scheduled_rate_advancing_through_a_skip():
-    """A schedule reads the clock the rule counts on, and the guard exempts clocks from the freeze, so a
-    skipped step still consumes a step and the next batch gets the next rate on the curve."""
+def test_a_guard_holds_the_rule_clock_back_through_a_skip():
+    """The rule's clock counts the updates the rule applied, and a skipped step is one it did not. A clock
+    that counted it anyway would bias-correct the next step against a moment that never saw it, so the
+    schedule reading that clock pauses for the step too."""
     p, loss = quadratic_problem()
     step = compile_train(loss, apply_if_finite(adam(cosine_schedule(0.1, 10))))
     clock = state_named(step, "adam/step_count")
@@ -120,10 +121,25 @@ def test_a_guard_keeps_a_scheduled_rate_advancing_through_a_skip():
     step(BAD)
     step(GOOD)
 
-    assert int(clock.get_value()) == 3
+    assert int(clock.get_value()) == 2
     assert np.all(np.isfinite(p.get_value()))
 
 
+def test_a_guard_leaves_a_caller_held_clock_counting_through_a_skip():
+    """A clock the caller holds is not the rule's state, so the guard does not touch it: it counts the
+    calls made, skipped or not, which is the other notion of time and the caller's to keep."""
+    p, loss = quadratic_problem()
+    clock = step_counter("batches")
+    step = compile_train(loss, apply_if_finite(adam(cosine_schedule(0.1, 10)(clock))))
+
+    step(GOOD)
+    step(BAD)
+    step(GOOD)
+
+    assert int(clock.get_value()) == 3
+    assert int(state_named(step, "adam/step_count").get_value()) == 2
+
+
 def test_weight_decay_still_reaches_a_step_the_guard_lets_through():
     """Decay is a transform inside the chain, so a guard wrapping the whole chain must leave it alone on
     a step that is applied."""
diff --git a/tests/optim/test_guards.py b/tests/optim/test_guards.py
index a8b30b2..2700ef5 100644
--- a/tests/optim/test_guards.py
+++ b/tests/optim/test_guards.py
@@ -50,16 +50,21 @@ def test_holds_back_parameters_and_state_on_a_nonfinite_step(poison):
     np.testing.assert_allclose(first_moment.get_value(), moment)
 
 
-def test_resumes_training_where_the_skipped_step_left_it():
+@pytest.mark.parametrize(
+    "make_rule", [lambda: sgd(0.1), lambda: adam(0.1)], ids=["stateless", "bias-corrected"]
+)
+def test_resumes_training_where_the_skipped_step_left_it(make_rule):
     """The skipped step must be a true no-op, not a damped or partial one: two good batches around it have
-    to land exactly where the same two would have landed on their own."""
+    to land exactly where the same two would have landed on their own. Adam is the case that can go
+    wrong quietly, since a step count that advanced through the skip would bias-correct the next step
+    against a moment that never saw it."""
     p, loss = poisonable_problem()
-    guarded = compile_train(loss, apply_if_finite(sgd(0.1)))
+    guarded = compile_train(loss, apply_if_finite(make_rule()))
     for batch in [GOOD, BAD, GOOD]:
         guarded(batch)
 
     reference_p, reference_loss = poisonable_problem()
-    unguarded = compile_train(reference_loss, sgd(0.1))
+    unguarded = compile_train(reference_loss, make_rule())
     for _ in range(2):
         unguarded(GOOD)
 
@@ -145,18 +150,6 @@ def test_skips_indefinitely_without_a_tolerance():
     np.testing.assert_allclose(p.get_value(), weights)
 
 
-def test_clocks_advance_through_a_skipped_step():
-    """A skipped step still consumed a step, and a frozen clock would stall every schedule reading it."""
-    p, loss = poisonable_problem()
-    step = compile_train(loss, apply_if_finite(adam(0.1)))
-    clock = state_named(step, "adam/step_count")
-
-    step(GOOD)
-    step(BAD)
-
-    assert int(clock.get_value()) == 2
-
-
 def test_large_step_holds_back_an_outsized_but_finite_step():
     """The case `nonfinite` cannot catch: a step far too big to be healthy, whose every value is finite."""
     p, loss = poisonable_problem()

From 0fc37c01aa94cfe281c8e7ef8485476ef5dbf3ac Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:31:00 +0800
Subject: [PATCH 08/16] Cast a numpy scalar rate to a Python float

---
 pytensor_ml/optim/base.py | 11 ++++++++---
 tests/optim/test_rules.py | 11 +++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 6ac39e0..331f604 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -253,8 +253,9 @@ def to_floatx(value: Rate) -> Rate:
 
     Parameters
     ----------
-    value : float or TensorVariable
-        A scalar a rule is about to build into its step.
+    value : float, numpy scalar, or TensorVariable
+        A scalar a rule is about to build into its step. A numpy scalar becomes a Python float, which
+        pytensor folds into the graph at the graph's own dtype, where a ``np.float64`` would widen it.
 
     Examples
     --------
@@ -267,7 +268,11 @@ def to_floatx(value: Rate) -> Rate:
 
         rate = to_floatx(1e-3)
     """
-    return value.astype(pytensor.config.floatX) if isinstance(value, Variable) else value
+    if isinstance(value, Variable):
+        return value.astype(pytensor.config.floatX)
+    if isinstance(value, np.generic | np.ndarray):
+        return float(value)
+    return value
 
 
 def get_gradients(
diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py
index 8151575..a7a11ba 100644
--- a/tests/optim/test_rules.py
+++ b/tests/optim/test_rules.py
@@ -627,3 +627,14 @@ def test_a_functional_rule_reads_a_schedule_off_its_own_clock():
 
     assert [clock.name for clock in clocks] == ["adam/step_count"]
     assert int(clocks[0].get_value()) == 2
+
+
+def test_a_numpy_scalar_rate_does_not_widen_a_float32_graph():
+    """A rate read back from a config or a checkpoint arrives as a numpy scalar, and a ``np.float64`` in a
+    float32 graph makes an update pytensor refuses with an error naming the parameter."""
+    with pytensor.config.change_flags(floatX="float32"):
+        p = params.trainable(np.zeros(3, dtype="float32"), name="w")
+        loss = (p**2).sum()
+        updates = sgd_updates(loss, [p], learning_rate=np.float64(0.1))
+
+        assert updates[p].dtype == "float32"

From 17baf9f70c0965dd617f606ece51635c57bca1aa Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:31:17 +0800
Subject: [PATCH 09/16] Pass sgd's namespace to the velocity it owns

---
 pytensor_ml/optim/alias.py |  2 +-
 tests/optim/test_rules.py  | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py
index 9fc7940..8a860ea 100644
--- a/pytensor_ml/optim/alias.py
+++ b/pytensor_ml/optim/alias.py
@@ -85,7 +85,7 @@ def rule(
             learning_rate=1.0,
             namespace=namespace,
         )
-        updates = trace(momentum, nesterov)(updates, parameters)
+        updates = trace(momentum, nesterov, namespace=namespace)(updates, parameters)
         return scale(learning_rate, namespace=namespace)(updates, parameters)
 
     return rule
diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py
index a7a11ba..9ff5c00 100644
--- a/tests/optim/test_rules.py
+++ b/tests/optim/test_rules.py
@@ -638,3 +638,13 @@ def test_a_numpy_scalar_rate_does_not_widen_a_float32_graph():
         updates = sgd_updates(loss, [p], learning_rate=np.float64(0.1))
 
         assert updates[p].dtype == "float32"
+
+
+def test_sgd_momentum_keeps_its_velocity_under_its_own_namespace():
+    """The ``namespace`` the docstring offers against colliding state has to reach the velocity too, or a
+    momentum sgd next to a ``trace`` still collides and the remedy on the sgd side does nothing."""
+    p = trainable(np.zeros(3), name="w")
+    loss = (p**2).sum()
+    updates = sgd(learning_rate=0.1, momentum=0.9, namespace="slow")(loss, [p])
+
+    assert {key.name for key in updates if key is not p} == {"w/slow/velocity"}

From 2ea0845ef970709889423fbb958e118b21812c2f Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:31:43 +0800
Subject: [PATCH 10/16] Keep guard and policy counters as integers

---
 pytensor_ml/optim/base.py   |  8 +++++---
 pytensor_ml/optim/guards.py |  4 ++--
 pytensor_ml/optim/policy.py | 14 +++++++-------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 331f604..6d8a8f5 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -567,9 +567,9 @@ def counter(name: str) -> Parameter:
     return step_counter(name)
 
 
-def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
+def scalar_state(name: str, fill_value: float = 0.0, dtype: str | None = None) -> Parameter:
     """
-    Allocate a floatX scalar shared variable.
+    Allocate a scalar shared variable, at ``floatX`` unless told otherwise.
 
     Parameters
     ----------
@@ -577,6 +577,8 @@ def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
         Name of the variable, used to match it at serialization boundaries.
     fill_value : float
         Value to initialize it with. Default 0.0.
+    dtype : str, optional
+        Storage dtype. Default ``floatX``. A count belongs in an integer dtype.
 
     Examples
     --------
@@ -589,7 +591,7 @@ def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
 
         scale = scalar_state("plateau/scale", fill_value=1.0)
     """
-    return pytensor.shared(np.asarray(fill_value, dtype=pytensor.config.floatX), name=name)
+    return pytensor.shared(np.asarray(fill_value, dtype=dtype or pytensor.config.floatX), name=name)
 
 
 def require_unique_state_names(updates: Updates) -> None:
diff --git a/pytensor_ml/optim/guards.py b/pytensor_ml/optim/guards.py
index aab79d4..ee7b415 100644
--- a/pytensor_ml/optim/guards.py
+++ b/pytensor_ml/optim/guards.py
@@ -201,7 +201,7 @@ def decide(updates: Updates, parameters: Sequence[Parameter]) -> TensorVariable:
 
 def _counter_or_new(given: Parameter | None, name: str) -> Parameter:
     """Return the caller's own counter, or a fresh scalar state slot under ``name``."""
-    return scalar_state(name) if given is None else given
+    return scalar_state(name, dtype="int64") if given is None else given
 
 
 def skip_if(
@@ -305,7 +305,7 @@ def guarded(
         consecutive = _counter_or_new(consecutive_skips, f"{namespace}/consecutive_skips")
         total = _counter_or_new(total_skips, f"{namespace}/total_skips")
 
-        next_consecutive = pt.where(skipping, consecutive + 1, 0.0)
+        next_consecutive = pt.where(skipping, consecutive + 1, 0)
         if max_consecutive_skips is not None:
             next_consecutive = CheckAndRaise(
                 FloatingPointError,
diff --git a/pytensor_ml/optim/policy.py b/pytensor_ml/optim/policy.py
index a725f4f..b4e2373 100644
--- a/pytensor_ml/optim/policy.py
+++ b/pytensor_ml/optim/policy.py
@@ -154,9 +154,9 @@ def policy(
         updates = Steps(result)
 
         best_loss = scalar_state(f"{namespace}/best_loss", fill_value=np.inf)
-        waited = scalar_state(f"{namespace}/wait")
-        cooling = scalar_state(f"{namespace}/cooldown")
-        observed = scalar_state(f"{namespace}/observed")
+        waited = scalar_state(f"{namespace}/wait", dtype="int64")
+        cooling = scalar_state(f"{namespace}/cooldown", dtype="int64")
+        observed = scalar_state(f"{namespace}/observed", dtype="int64")
         mean_loss = scalar_state(f"{namespace}/mean_loss")
 
         # Everything below is gated on `deciding`, so a window that is still filling advances nothing. At the
@@ -166,23 +166,23 @@ def policy(
         deciding = seen >= accumulation_size
 
         improved = deciding & (running_mean < (1 - rtol) * best_loss - atol)
-        counted = pt.where(deciding, pt.where(improved, 0.0, waited + 1), waited)
+        counted = pt.where(deciding, pt.where(improved, 0, waited + 1), waited)
 
         # Cooling down zeroes the count rather than pausing it, so the steps immediately after a cut cannot
         # add up to the next one before the network has had a chance to respond to the rate it just got.
         in_cooldown = cooling > 0
         cutting = deciding & ~in_cooldown & (counted >= patience)
-        next_cooling = pt.where(in_cooldown, cooling - 1, pt.where(cutting, cooldown, 0.0))
+        next_cooling = pt.where(in_cooldown, cooling - 1, pt.where(cutting, cooldown, 0))
 
         updates[scale] = pt.where(cutting, pt.maximum(scale * factor, min_scale), scale).astype(
             scale.dtype
         )
         updates[best_loss] = pt.where(improved, running_mean, best_loss).astype(best_loss.dtype)
-        updates[waited] = pt.where(deciding & (in_cooldown | cutting), 0.0, counted).astype(
+        updates[waited] = pt.where(deciding & (in_cooldown | cutting), 0, counted).astype(
             waited.dtype
         )
         updates[cooling] = pt.where(deciding, next_cooling, cooling).astype(cooling.dtype)
-        updates[observed] = pt.where(deciding, 0.0, seen).astype(observed.dtype)
+        updates[observed] = pt.where(deciding, 0, seen).astype(observed.dtype)
         updates[mean_loss] = pt.where(deciding, 0.0, running_mean).astype(mean_loss.dtype)
 
         return updates

From 052c3156ea44c35e67bb6088ceceb4ca4903bfad Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:32:07 +0800
Subject: [PATCH 11/16] Delete counter in favor of step_counter

---
 docs/source/api/optim.rst          |  1 -
 pytensor_ml/optim/__init__.py      |  2 --
 pytensor_ml/optim/base.py          | 42 ++++--------------------------
 pytensor_ml/optim/rules.py         |  8 +++---
 tests/optim/test_public_surface.py |  2 +-
 5 files changed, 10 insertions(+), 45 deletions(-)

diff --git a/docs/source/api/optim.rst b/docs/source/api/optim.rst
index 2391691..632a92c 100644
--- a/docs/source/api/optim.rst
+++ b/docs/source/api/optim.rst
@@ -86,7 +86,6 @@ Building blocks
 .. autosummary::
     :toctree: generated/
 
-    counter
     get_gradients
     scalar_state
     state_for
diff --git a/pytensor_ml/optim/__init__.py b/pytensor_ml/optim/__init__.py
index 5f0a595..a86eeb0 100644
--- a/pytensor_ml/optim/__init__.py
+++ b/pytensor_ml/optim/__init__.py
@@ -18,7 +18,6 @@
     Transform,
     Updates,
     chain,
-    counter,
     get_gradients,
     scalar_state,
     state_for,
@@ -93,7 +92,6 @@
     "compile_train",
     "constant_schedule",
     "cosine_schedule",
-    "counter",
     "exponential_schedule",
     "get_gradients",
     "join_schedules",
diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 6d8a8f5..3d2c963 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -10,7 +10,7 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.sharedvar import TensorSharedVariable
 
-from pytensor_ml.params import TrainableParameter, step_counter
+from pytensor_ml.params import StepCounter, TrainableParameter, step_counter
 from pytensor_ml.pytensorf import rewrite_pregrad
 
 type Parameter = TensorSharedVariable
@@ -191,7 +191,7 @@ def inverse_square_root(step_count):
 """
 
 
-def rate_on(learning_rate: LearningRate, clock: Parameter) -> Rate:
+def rate_on(learning_rate: LearningRate, clock: StepCounter) -> Rate:
     """
     Read a schedule off ``clock``. Any other rate passes through untouched.
 
@@ -199,7 +199,7 @@ def rate_on(learning_rate: LearningRate, clock: Parameter) -> Rate:
     ----------
     learning_rate : LearningRate
         A rate, or a schedule of the step count.
-    clock : shared tensor variable
+    clock : StepCounter
         The training clock a schedule is evaluated at, ordinarily the one the rule counts its own steps on.
 
     Returns
@@ -234,8 +234,8 @@ def read_rate(learning_rate: LearningRate, namespace: str) -> tuple[Rate, Update
     """
     if not callable(learning_rate):
         return to_floatx(learning_rate), Updates()
-    clock = counter(f"{namespace}/step_count")
-    return to_floatx(learning_rate(clock)), Updates({clock: clock + 1})
+    clock = step_counter(f"{namespace}/step_count")
+    return to_floatx(learning_rate(clock)), Updates({clock: clock.advance()})
 
 
 def to_floatx(value: Rate) -> Rate:
@@ -535,38 +535,6 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param
     return state
 
 
-def counter(name: str) -> Parameter:
-    """
-    Allocate the training clock a component counts its own steps on.
-
-    A :class:`~pytensor_ml.params.StepCounter` rather than a plain shared variable, so a schedule can read
-    the same notion of time the rule uses, and :func:`~pytensor_ml.pytensorf.collect_clock_updates` advances
-    it for a caller who does not write the advance themselves.
-
-    Parameters
-    ----------
-    name : str
-        Name of the clock, used to match it at serialization boundaries.
-
-    Returns
-    -------
-    clock : StepCounter
-        A new step counter under ``name``.
-
-    Examples
-    --------
-    Read a schedule off a clock of your own, which is what a transform does when it applies a rate after
-    the rule rather than inside it:
-
-    .. code-block:: python
-
-        from pytensor_ml.optim import cosine_schedule, counter
-
-        rate = cosine_schedule(3e-4, total_steps=10_000)(counter("my_transform/step_count"))
-    """
-    return step_counter(name)
-
-
 def scalar_state(name: str, fill_value: float = 0.0, dtype: str | None = None) -> Parameter:
     """
     Allocate a scalar shared variable, at ``floatX`` unless told otherwise.
diff --git a/pytensor_ml/optim/rules.py b/pytensor_ml/optim/rules.py
index e10f80b..c6a6b60 100644
--- a/pytensor_ml/optim/rules.py
+++ b/pytensor_ml/optim/rules.py
@@ -13,13 +13,13 @@
     Parameter,
     Steps,
     Updates,
-    counter,
     gradients_to_descend,
     rate_on,
     read_rate,
     state_for,
     to_floatx,
 )
+from pytensor_ml.params import step_counter
 
 
 def sgd_updates(
@@ -190,7 +190,7 @@ def _adam_family_updates(
         Predicate selecting which parameters the decay reaches. Every parameter when omitted.
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    step_count = counter(f"{namespace}/step_count")
+    step_count = step_counter(f"{namespace}/step_count")
     learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
@@ -408,7 +408,7 @@ def nadam_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    step_count = counter(f"{namespace}/step_count")
+    step_count = step_counter(f"{namespace}/step_count")
     learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
@@ -504,7 +504,7 @@ def adamax_updates(
         loss_value = step(np.zeros((8, 4)), np.zeros((8, 1)))
     """
     incoming, gradients = gradients_to_descend(loss_gradients_or_updates, parameters, namespace)
-    step_count = counter(f"{namespace}/step_count")
+    step_count = step_counter(f"{namespace}/step_count")
     learning_rate = to_floatx(rate_on(learning_rate, step_count))
     new_step_count = step_count + 1
     new_step_count_float = new_step_count.astype(config.floatX)
diff --git a/tests/optim/test_public_surface.py b/tests/optim/test_public_surface.py
index 1fcda86..f142c4c 100644
--- a/tests/optim/test_public_surface.py
+++ b/tests/optim/test_public_surface.py
@@ -16,7 +16,7 @@ def test_exported_names_are_unique_and_sorted():
     assert optim.__all__ == sorted(set(optim.__all__))
 
 
-@pytest.mark.parametrize("name", ["to_updates", "steps_of", "state_for", "counter", "chain"])
+@pytest.mark.parametrize("name", ["to_updates", "steps_of", "state_for", "scalar_state", "chain"])
 def test_the_transform_authoring_api_stays_public(name):
     """The ``Transform`` docstring tells authors to write a transform with these. Dropping one from
     ``__all__`` would leave the documented path running through ``pytensor_ml.optim.base``."""

From 0f1a078b07711157ec46197191a3cb8c66cd2960 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:32:43 +0800
Subject: [PATCH 12/16] Refuse a plateau scale the rule never reads

---
 pytensor_ml/optim/policy.py | 13 ++++++++++++-
 tests/optim/test_policy.py  | 12 ++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/pytensor_ml/optim/policy.py b/pytensor_ml/optim/policy.py
index b4e2373..e1de025 100644
--- a/pytensor_ml/optim/policy.py
+++ b/pytensor_ml/optim/policy.py
@@ -3,6 +3,7 @@
 import numpy as np
 import pytensor.tensor as pt
 
+from pytensor.graph import graph_inputs
 from pytensor.tensor import TensorVariable
 
 from pytensor_ml.optim.base import (
@@ -57,7 +58,9 @@ def reduce_on_plateau(
     rule : Transform
         The rule to wrap. Its rate must be built from ``scale`` for the cuts to reach the step.
     scale : shared tensor variable
-        The multiplier this policy owns. Nothing else may write it.
+        The multiplier this policy owns. Nothing else may write it, and the rule's step must read it, or
+        the policy raises when invoked. Two invocations of one policy each keep a fresh history but cut
+        this same variable, so two training functions built from one policy cut it twice as fast.
     factor : float
         Multiplier applied on a cut, in the open interval (0, 1). Default 0.1.
     patience : int or TensorVariable
@@ -152,6 +155,14 @@ def policy(
                 "would move uphill. Put an optimizer such as `adam(rate)` inside the policy."
             )
         updates = Steps(result)
+        steps = [updates[parameter] for parameter in parameters]
+        if steps and scale not in graph_inputs(steps):
+            raise ValueError(
+                f"reduce_on_plateau would cut {scale.name!r}, but no parameter's step reads it, so the "
+                "cuts would reach nothing. Build the rule's rate from this same variable, e.g. "
+                "`adam(learning_rate=scale * 1e-3)`, and pass that one object here. Two `scalar_state` "
+                "calls with one name are two variables."
+            )
 
         best_loss = scalar_state(f"{namespace}/best_loss", fill_value=np.inf)
         waited = scalar_state(f"{namespace}/wait", dtype="int64")
diff --git a/tests/optim/test_policy.py b/tests/optim/test_policy.py
index 121a0a7..b9cf879 100644
--- a/tests/optim/test_policy.py
+++ b/tests/optim/test_policy.py
@@ -127,6 +127,18 @@ def noisy_run(cooldown, min_scale, n_steps=120):
     assert guarded_floor == pytest.approx(0.01, rel=RTOL)
 
 
+def test_rejects_a_scale_the_rule_does_not_read():
+    """Cutting a multiplier no step reads is a policy that looks configured and does nothing. The usual way
+    to get there is two ``scalar_state`` calls under one name, which are two variables."""
+    p, loss = plateaued_problem()
+    rate_scale = scalar_state("plateau/scale", fill_value=1.0)
+    policy_scale = scalar_state("plateau/scale", fill_value=1.0)
+    rule = reduce_on_plateau(adam(learning_rate=rate_scale * 0.05), policy_scale)
+
+    with pytest.raises(ValueError, match="no parameter's step reads it"):
+        rule(loss, [p])
+
+
 def test_rejects_precomputed_gradients():
     """The policy decides from the loss, and `loss_or_gradients` is a union, so a caller passing gradients
     would otherwise hand it a list and get a confusing failure deep in the comparison."""

From 702d71eb37c9a26fbd46d6f628fa6612a5d272a1 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:46:41 +0800
Subject: [PATCH 13/16] Rename the cloudpickle test to drop its stale caveat

---
 tests/optim/test_composition.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/optim/test_composition.py b/tests/optim/test_composition.py
index dd9ad7f..d58032a 100644
--- a/tests/optim/test_composition.py
+++ b/tests/optim/test_composition.py
@@ -208,7 +208,7 @@ def test_two_scheduled_scales_sharing_a_namespace_collide_loudly():
     compile_train(loss, apart)
 
 
-def test_an_unused_rule_ships_through_cloudpickle_and_trains_after():
+def test_a_configured_rule_ships_through_cloudpickle_and_trains_after():
     """A fit config that carries a configured rule to a remote worker goes through cloudpickle, which
     serializes a nested function's globals by value. A rule holds hyperparameters and nothing else, so the
     whole spine ships, and the state it allocates on the worker is the worker's."""

From 1c2b3e609d783aba0256232a7bb8efd202c1a683 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:46:41 +0800
Subject: [PATCH 14/16] Assert exact first steps behind a gradient-space
 transform

---
 tests/optim/test_training.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py
index b4d9b32..4261f1c 100644
--- a/tests/optim/test_training.py
+++ b/tests/optim/test_training.py
@@ -269,8 +269,9 @@ def per_group(loss_or_gradients, parameters):
     )
     step()
 
-    assert weight.get_value() < 2.0
-    assert bias.get_value() < 3.0
+    # The clip's bound is far above the gradient norm, so both first steps are adam's exact sign descent.
+    np.testing.assert_allclose(weight.get_value(), [2.0 - 0.1 * (1 + 0.01 * 2.0)], rtol=1e-3)
+    np.testing.assert_allclose(bias.get_value(), [3.0 - 0.1], rtol=1e-3)
 
 
 def test_an_updates_dict_missing_a_collected_parameter_is_refused():

From f5ea881dbaa5553cc50c36be16474dc9d19176e8 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 12:46:42 +0800
Subject: [PATCH 15/16] Cover every clock-writing rule in the checkpoint clock
 test

---
 tests/test_checkpoint.py | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
index 9cdbdf0..70fe7f3 100644
--- a/tests/test_checkpoint.py
+++ b/tests/test_checkpoint.py
@@ -12,7 +12,7 @@
 from pytensor_ml.checkpoint import generator_from_state, load_state, save_state
 from pytensor_ml.layers import BatchNorm, Dropout, Linear, Sequential
 from pytensor_ml.loss import SquaredError, supervised_loss
-from pytensor_ml.optim import adam, cosine_schedule, sgd
+from pytensor_ml.optim import adadelta, adagrad, adam, cosine_schedule, rmsprop, sgd
 from pytensor_ml.params import trainable
 from pytensor_ml.pytensorf import (
     collect_optimizer_state,
@@ -633,14 +633,29 @@ def test_collect_optimizer_state_is_empty_for_a_rule_that_keeps_none():
     assert collect_optimizer_state(sgd(learning_rate=1e-2)(loss, parameters), parameters) == []
 
 
-def test_a_scheduled_rule_puts_the_clock_it_reads_in_its_updates():
+@pytest.mark.parametrize(
+    "alias, clock_name",
+    [
+        (lambda schedule: sgd(learning_rate=schedule), "sgd/step_count"),
+        (lambda schedule: sgd(learning_rate=schedule, momentum=0.9), "sgd/step_count"),
+        (lambda schedule: rmsprop(learning_rate=schedule), "rmsprop/step_count"),
+        (lambda schedule: adagrad(learning_rate=schedule), "adagrad/step_count"),
+        (lambda schedule: adadelta(learning_rate=schedule), "adadelta/step_count"),
+    ],
+    ids=["sgd", "sgd-momentum", "rmsprop", "adagrad", "adadelta"],
+)
+def test_a_scheduled_rule_puts_the_clock_it_reads_in_its_updates(alias, clock_name):
     """A schedule's clock is state the rule carries between steps, so it has to be in the updates dict a
     checkpoint is taken from. A rule that read the clock without writing it would train correctly and
-    still hand a checkpoint that restarts the schedule at zero."""
+    still hand a checkpoint that restarts the schedule at zero. These are the rules with no step count of
+    their own, and momentum sgd reads the schedule through ``scale``, a separate path."""
     X = pt.tensor("X", shape=(None, 4))
     loss, target = supervised_loss(Linear("fc", n_in=4, n_out=1)(X), SquaredError())
     parameters = collect_trainable_params(loss)
 
-    updates = sgd(learning_rate=cosine_schedule(1e-2, total_steps=10))(loss, parameters)
+    updates = alias(cosine_schedule(1e-2, total_steps=10))(loss, parameters)
+    clocks = [
+        v.name for v in collect_optimizer_state(updates, parameters) if "step_count" in v.name
+    ]
 
-    assert [v.name for v in collect_optimizer_state(updates, parameters)] == ["sgd/step_count"]
+    assert clocks == [clock_name]

From 9bf46953c36be32b212abda11d90c147e2ea0f31 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Fri, 18 Sep 2026 22:18:07 +0800
Subject: [PATCH 16/16] Drop the scale_by_schedule test that duplicates the
 scale schedule test

---
 tests/optim/test_transform.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/tests/optim/test_transform.py b/tests/optim/test_transform.py
index 020769b..cb7e97b 100644
--- a/tests/optim/test_transform.py
+++ b/tests/optim/test_transform.py
@@ -22,22 +22,6 @@ def test_scale_applies_factor():
     np.testing.assert_allclose(function([], out[p])(), [0.5, -1.0])
 
 
-def test_scale_by_schedule_applies_the_rate_its_clock_reads():
-    p = trainable(np.zeros(2), name="w")
-    updates = {p: p + pt.constant(np.array([2.0, -4.0]))}
-    out = scale_by_schedule(linear_schedule(1.0, total_steps=4, final_learning_rate=0.0))(
-        updates, [p]
-    )
-
-    (clock,) = collect_step_counters(out[p])
-    step = function([], out[p], updates={clock: clock.advance()})
-
-    # The clock starts at zero, where the schedule is still at its initial rate of 1.0.
-    np.testing.assert_allclose(step(), [2.0, -4.0])
-    # One quarter of the horizon later the rate is 0.75.
-    np.testing.assert_allclose(step(), [1.5, -3.0])
-
-
 def test_scale_by_schedule_allocates_a_clock_per_namespace():
     """Two scheduled scales in one graph measure their own time. Sharing a clock by default would make the
     second one's schedule start wherever the first had already advanced it to."""