From be6b9bdd1b0f141099bab442e0f54ffae6e262e3 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Sun, 9 Aug 2026 18:03:58 -0500
Subject: [PATCH 1/5] Count rule steps on a training clock

compile_train accepts a rule's own advance rather than rejecting it, so a rule's updates keep advancing their counter when compiled directly instead of through compile_train.
---
 pytensor_ml/optim/base.py         | 14 ++++---
 pytensor_ml/optim/train.py        | 16 +++++---
 pytensor_ml/pytensorf/collect.py  |  5 ++-
 tests/optim/test_learning_rate.py | 66 +++++++++++++++++++++++++++++--
 tests/optim/test_rules.py         | 15 +++++++
 5 files changed, 100 insertions(+), 16 deletions(-)

diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py
index 02d9732..6be55b5 100644
--- a/pytensor_ml/optim/base.py
+++ b/pytensor_ml/optim/base.py
@@ -10,6 +10,7 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.sharedvar import TensorSharedVariable
 
+from pytensor_ml.params import step_counter
 from pytensor_ml.pytensorf import rewrite_pregrad
 
 type Parameter = TensorSharedVariable
@@ -150,11 +151,14 @@ def allocate() -> Parameter:
 
 
 def counter(name: str) -> Parameter:
-    """Return an int64 step counter shared variable initialized to zero, reused across invocations of a
-    rule wrapped in :func:`reuses_state` so its step count keeps advancing."""
-    return _reuse_or_allocate(
-        name, lambda: pytensor.shared(np.asarray(0, dtype="int64"), name=name)
-    )
+    """Return the training clock a rule 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
+    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.
+    """
+    return _reuse_or_allocate(name, lambda: step_counter(name))
 
 
 def scalar_state(name: str, fill_value: float = 0.0) -> Parameter:
diff --git a/pytensor_ml/optim/train.py b/pytensor_ml/optim/train.py
index 4e855a5..386a8fc 100644
--- a/pytensor_ml/optim/train.py
+++ b/pytensor_ml/optim/train.py
@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 
 from pytensor.compile import Function, SharedVariable
-from pytensor.graph.basic import Variable
+from pytensor.graph.basic import Variable, equal_computations
 from pytensor.tensor import TensorVariable
 
 from pytensor_ml.optim.base import Parameter, UpdateRule, require_unique_state_names
@@ -79,12 +79,18 @@ def compile_train(
     for clock, next_count in collect_clock_updates(
         [loss, *extra_outputs, *updates.values()]
     ).items():
-        if clock in updates:
+        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 updated by the rule. Clocks advance once "
-                "per step on their own; drop the explicit update."
+                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
 
     require_unique_state_names(updates)
 
diff --git a/pytensor_ml/pytensorf/collect.py b/pytensor_ml/pytensorf/collect.py
index fe86cec..7ba446a 100644
--- a/pytensor_ml/pytensorf/collect.py
+++ b/pytensor_ml/pytensorf/collect.py
@@ -122,8 +122,9 @@ def collect_clock_updates(
     if len(step_counts) > 1:
         raise ValueError(
             f"Training clocks {sorted(str(counter.name) for counter in counters)} hold different step "
-            f"counts {sorted(step_counts)}. They all count training steps, so this means some of them were "
-            "restored from a checkpoint and others were not."
+            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}
 
diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index 1c3cb02..4c79c63 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -16,7 +16,9 @@
     sgd,
     substitute_schedule,
 )
-from pytensor_ml.params import step_counter, trainable
+from pytensor_ml.params import StepCounter, step_counter, trainable
+from pytensor_ml.pytensorf import collect_step_counters
+from pytensor_ml.pytensorf import function as pytensor_function
 
 RTOL = 1e-6
 
@@ -213,8 +215,9 @@ def test_two_training_functions_share_a_held_clock():
     assert int(clock.get_value()) == 6
 
 
-def test_compile_train_rejects_a_rule_that_advances_the_clock_itself():
-    """Two writers would leave the clock holding one advance while both were configured."""
+def test_compile_train_accepts_a_rule_that_advances_the_clock_the_same_way():
+    """Rules advance the clock they count their own steps on, so their write and the collected one agree.
+    Accepting it is what lets a rule's updates stand on their own outside compile_train."""
     _, loss = quadratic_problem()
     clock = step_counter()
 
@@ -222,7 +225,23 @@ def rule(loss_or_gradients, parameters):
         updates = sgd(learning_rate=cosine_schedule(0.1, 10)(clock))(loss_or_gradients, parameters)
         return {**updates, clock: clock + 1}
 
-    with pytest.raises(ValueError, match="already updated by the rule"):
+    step = compile_train(loss, rule)
+    step()
+    step()
+
+    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)
 
 
@@ -281,3 +300,42 @@ def test_a_clock_that_shadows_a_parameter_name_is_caught():
 
     with pytest.raises(ValueError, match="share the name 'w'"):
         compile_train(loss, sgd(learning_rate=cosine_schedule(0.1, 10)(shadowing_clock)))
+
+
+def test_a_rule_counts_its_own_steps_on_a_clock():
+    """The counter adam keeps for bias correction is a clock, so it is collected as one and a schedule can
+    read the same notion of time the rule uses."""
+    p, loss = quadratic_problem()
+    updates = adam(learning_rate=0.1)(loss, [p])
+
+    counters = collect_step_counters(list(updates.values()))
+    assert [counter.name for counter in counters] == ["adam/step_count"]
+    assert all(isinstance(counter, StepCounter) for counter in counters)
+
+
+def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
+    """Two clocks can still coexist when the caller holds one of their own, and a checkpoint that restored
+    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")
+    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")])
+
+
+def test_a_caller_held_clock_keeps_advancing_for_other_readers():
+    """A clock is state that outlives the compile: the caller reads it, and so can another compiled function.
+    Compiling a training step must not take it out of the set that step advances."""
+    p, loss = quadratic_problem()
+    clock = step_counter("agent_clock")
+    step = compile_train(loss, adam(learning_rate=cosine_schedule(0.1, 100)(clock)))
+    epsilon = pytensor_function([], cosine_schedule(1.0, 100)(clock))
+
+    for _ in range(20):
+        step()
+
+    assert int(clock.get_value()) == 20
+    assert float(epsilon()) < 1.0
diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py
index 3eba103..2419982 100644
--- a/tests/optim/test_rules.py
+++ b/tests/optim/test_rules.py
@@ -152,6 +152,21 @@ def test_adam_first_step_is_sign_descent():
     np.testing.assert_allclose(step, lr * np.sign(start), rtol=RTOL)
 
 
+def test_rule_updates_advance_their_counter_without_compile_train():
+    """adam_updates has to stand on its own: bias correction reads the counter, so a caller who compiles the
+    updates directly must get an advancing count rather than one frozen at the first step."""
+    p = trainable(np.zeros(3), name="w")
+    loss = (p**2).sum()
+    updates = adam_updates(loss, [p])
+    clock = next(key for key in updates if key.name == "adam/step_count")
+
+    step = function([], loss, updates=updates)
+    step()
+    step()
+
+    assert int(clock.get_value()) == 2
+
+
 def test_adam_updates_keyed_by_object_with_named_state():
     """State is discovered by object identity; names exist only for serialization."""
     p = trainable(np.zeros(3), name="w")

From 94fbe0f0fa963ec59b2217295db1f7b38298e18d Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Sun, 9 Aug 2026 18:04:57 -0500
Subject: [PATCH 2/5] Read schedules off the clock the rule counts on

The applied rate is an expression over the clock now, so there is no rate variable to publish it to; ask for it as an extra output instead.
---
 pytensor_ml/optim/__init__.py     |   4 -
 pytensor_ml/optim/alias.py        |  16 ++--
 pytensor_ml/optim/rules.py        |   2 +-
 pytensor_ml/optim/schedules.py    |  35 +++++----
 pytensor_ml/optim/transform.py    | 122 +-----------------------------
 tests/optim/test_learning_rate.py | 102 +++++++++++++++----------
 tests/optim/test_schedules.py     |  15 ++--
 tests/optim/test_transform.py     | 116 ++--------------------------
 8 files changed, 106 insertions(+), 306 deletions(-)

diff --git a/pytensor_ml/optim/__init__.py b/pytensor_ml/optim/__init__.py
index b35d33e..1621412 100644
--- a/pytensor_ml/optim/__init__.py
+++ b/pytensor_ml/optim/__init__.py
@@ -45,8 +45,6 @@
 from pytensor_ml.optim.transform import (
     add_weight_decay,
     scale,
-    scale_by_schedule,
-    substitute_schedule,
     trace,
 )
 
@@ -87,10 +85,8 @@
     "rprop_updates",
     "scalar_state",
     "scale",
-    "scale_by_schedule",
     "sgd",
     "sgd_updates",
     "step_decay",
-    "substitute_schedule",
     "trace",
 ]
diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py
index 8fd42b7..6f70e4d 100644
--- a/pytensor_ml/optim/alias.py
+++ b/pytensor_ml/optim/alias.py
@@ -7,8 +7,8 @@
     Rate,
     UpdateRule,
     Updates,
+    counter,
     reuses_state,
-    scalar_state,
 )
 from pytensor_ml.optim.rules import (
     _require_numeric_learning_rate,
@@ -22,7 +22,7 @@
     rprop_updates,
     sgd_updates,
 )
-from pytensor_ml.optim.transform import scale, substitute_schedule, trace
+from pytensor_ml.optim.transform import scale, trace
 
 
 def _at_learning_rate(
@@ -30,13 +30,13 @@ def _at_learning_rate(
     name: str,
     build_updates: Callable[[Rate], Updates],
 ) -> Updates:
-    """Build updates at ``learning_rate``, substituting a schedule into a shared rate variable named
-    ``"{name}/learning_rate"`` when one was given."""
+    """Build updates at ``learning_rate``, reading a schedule off the rule's step clock. A rule that keeps its
+    own clock under another name ends up with two, which is harmless: every clock advances once per step from
+    zero, so they agree about the step at every point."""
     if not callable(learning_rate):
         return build_updates(learning_rate)
 
-    rate_variable = scalar_state(f"{name}/learning_rate")
-    return substitute_schedule(build_updates(rate_variable), rate_variable, learning_rate)
+    return build_updates(learning_rate(counter(f"{name}/step_count")))
 
 
 def sgd(
@@ -50,8 +50,8 @@ def sgd(
     learning_rate : float, shared tensor variable, symbolic scalar, or Schedule
         Step size. A float is baked into the graph; a scalar shared variable can be steered from Python with
         ``set_value``; any scalar graph is used as the rate directly, as in
-        ``cosine_schedule(3e-4, 10_000)(step_counter())``; an unapplied schedule is substituted into the
-        rate the rule uses, on-graph, from an owned step counter. Default 0.01.
+        ``cosine_schedule(3e-4, 10_000)(step_counter())``; an unapplied schedule is read off the clock the
+        rule counts its own steps on. Default 0.01.
     momentum : float
         Momentum coefficient. A value of 0 (the default) gives plain SGD.
     nesterov : bool
diff --git a/pytensor_ml/optim/rules.py b/pytensor_ml/optim/rules.py
index c14450d..63d94a0 100644
--- a/pytensor_ml/optim/rules.py
+++ b/pytensor_ml/optim/rules.py
@@ -549,7 +549,7 @@ def _require_numeric_learning_rate(learning_rate: LearningRate) -> None:
             "rprop's learning rate initializes its per-parameter step sizes rather than scaling the step, "
             "so it must be a plain number and cannot be scheduled or steered. Schedule a rule whose rate "
             "multiplies the step, or scale rprop's finished step with "
-            "`chain(rprop(...), scale_by_schedule(...))`, which is a different algorithm."
+            "`chain(rprop(...), scale(curve(clock)))`, which is a different algorithm."
         )
 
 
diff --git a/pytensor_ml/optim/schedules.py b/pytensor_ml/optim/schedules.py
index 819623c..6b424cd 100644
--- a/pytensor_ml/optim/schedules.py
+++ b/pytensor_ml/optim/schedules.py
@@ -61,8 +61,9 @@ def cosine_schedule(
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
 
     Examples
     --------
@@ -125,8 +126,9 @@ def linear_schedule(
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
     _validate_horizon(total_steps, transition_begin)
 
@@ -179,8 +181,9 @@ def exponential_schedule(
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
     _validate_horizon(total_steps, transition_begin)
     if final_learning_rate <= 0.0 or learning_rate <= 0.0:
@@ -242,8 +245,9 @@ def polynomial_schedule(
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
     _validate_horizon(total_steps, transition_begin)
     if power <= 0.0:
@@ -300,8 +304,9 @@ def step_decay(
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
     if decay_every < 1:
         raise ValueError(f"decay_every must be at least 1, got {decay_every}.")
@@ -337,8 +342,9 @@ def constant_schedule(learning_rate: float) -> Schedule:
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
 
     def schedule(step_count: TensorVariable) -> TensorVariable:
@@ -374,8 +380,9 @@ def join_schedules(schedules: Sequence[Schedule], boundaries: Sequence[int]) ->
     Returns
     -------
     Schedule
-        A callable mapping the symbolic step count to a scalar learning rate, for
-        :func:`~pytensor_ml.optim.transform.scale_by_schedule`.
+        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
+        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
+        ``learning_rate`` to have the rule read it off the clock it already counts on.
     """
     if not schedules:
         raise ValueError("join_schedules needs at least one schedule.")
diff --git a/pytensor_ml/optim/transform.py b/pytensor_ml/optim/transform.py
index 7022ab3..2393cbf 100644
--- a/pytensor_ml/optim/transform.py
+++ b/pytensor_ml/optim/transform.py
@@ -1,20 +1,6 @@
 from collections.abc import Callable, Sequence
 
-import pytensor.tensor as pt
-
-from pytensor.graph.basic import Variable
-from pytensor.graph.replace import graph_replace
-
-from pytensor_ml.optim.base import (
-    Parameter,
-    Rate,
-    Schedule,
-    Transform,
-    Updates,
-    counter,
-    scalar_state,
-    state_for,
-)
+from pytensor_ml.optim.base import Parameter, Rate, Transform, Updates, state_for
 
 
 def trace(decay: float = 0.9, nesterov: bool = False) -> Transform:
@@ -79,112 +65,6 @@ def transform(updates: Updates, parameters: Sequence[Parameter]) -> Updates:
     return transform
 
 
-def scale_by_schedule(schedule_fn: Schedule, learning_rate: Parameter | None = None) -> Transform:
-    """
-    Scale each step by a learning rate produced from an owned step counter.
-
-    The counter starts at zero and increments by one each call. ``schedule_fn`` receives it symbolically and
-    must return a scalar learning rate, so the schedule lives on the graph and stays synchronized with the
-    optimizer's own time step.
-
-    The rate *this transform* applies is also written to a scalar shared variable, so it can be checkpointed
-    alongside the counter and read from Python after a step. Pass ``learning_rate`` to supply that variable
-    and hold a reference to it; otherwise one is allocated as ``"schedule/learning_rate"``. Two scheduled
-    scalings in one chain each publish their own factor and share one step counter, so the rate reaching the
-    parameters is their product. The graph writes the variable on every step and never reads it, so
-    ``set_value`` on it does not steer a schedule that is a function of the step count.
-
-    Parameters
-    ----------
-    schedule_fn : callable
-        Maps the symbolic step counter (a TensorVariable) to a scalar learning rate.
-    learning_rate : shared tensor variable, optional
-        Scalar variable to publish the applied rate to, cast to its dtype. Allocated internally when
-        omitted.
-
-    Returns
-    -------
-    Transform
-        A transform that rescales the updates dict by the scheduled rate, publishes it, and advances the
-        counter.
-    """
-
-    def transform(updates: Updates, parameters: Sequence[Parameter]) -> Updates:
-        step_count = counter("schedule/step_count")
-        published_rate = (
-            scalar_state("schedule/learning_rate") if learning_rate is None else learning_rate
-        )
-        if published_rate in updates:
-            raise ValueError(
-                f"Two scheduled scalings in one chain would share the state {published_rate.name!r}, "
-                "leaving it holding only the second rate while both are applied. Multiply the schedules "
-                "into a single schedule function, or pass distinct `learning_rate` variables."
-            )
-
-        rate = pt.cast(schedule_fn(step_count), published_rate.type.dtype)
-
-        next_updates = dict(updates)
-        for parameter in parameters:
-            next_updates[parameter] = parameter + rate * (updates[parameter] - parameter)
-        next_updates[published_rate] = rate
-        next_updates[step_count] = step_count + 1
-        return next_updates
-
-    return transform
-
-
-def substitute_schedule(
-    updates: Updates, learning_rate: Parameter, schedule_fn: Schedule
-) -> Updates:
-    """
-    Replace ``learning_rate`` with a scheduled rate everywhere it appears in ``updates``.
-
-    Where :func:`scale_by_schedule` multiplies the finished step, this substitutes the rate into the graph
-    at the point the rule uses it. The two agree whenever the rate is a plain multiplier on the step, and
-    only substitution is correct when it is not — a rule whose accumulators consume the rate, or a decay
-    term deliberately left unscaled.
-
-    The scheduled rate is also written back to ``learning_rate``, so the variable holds the rate that was
-    applied and can be read from Python or checkpointed. Raise if ``learning_rate`` never appears in
-    ``updates``, which means the rule does not take its rate from the graph and cannot be scheduled this
-    way.
-
-    Parameters
-    ----------
-    updates : Updates
-        Updates built against ``learning_rate`` as a symbolic rate.
-    learning_rate : shared tensor variable
-        Scalar variable the rule read its rate from.
-    schedule_fn : callable
-        Maps the symbolic step counter to a scalar learning rate.
-
-    Returns
-    -------
-    Updates
-        The updates with the scheduled rate substituted in, the rate published, and the counter advanced.
-    """
-    step_count = counter("schedule/step_count")
-    rate = pt.cast(schedule_fn(step_count), learning_rate.type.dtype)
-
-    # graph_replace rejects a replacement it cannot use, and the counter update contains no rate, so the
-    # whole output list goes in one call rather than one output at a time.
-    replacements: dict[Variable, Variable] = {learning_rate: rate}
-    update_keys = list(updates)
-    try:
-        replaced = graph_replace([updates[key] for key in update_keys], replacements)
-    except ValueError as error:
-        raise ValueError(
-            f"The learning rate {learning_rate.name!r} does not appear in the updates it was built with, "
-            "so a schedule cannot be substituted into them. This happens when a rule uses its rate for "
-            "something other than the graph, such as initializing per-parameter state."
-        ) from error
-
-    next_updates: Updates = dict(zip(update_keys, replaced))  # type: ignore[arg-type]
-    next_updates[learning_rate] = rate
-    next_updates[step_count] = step_count + 1
-    return next_updates
-
-
 def add_weight_decay(
     weight_decay: float = 0.01,
     mask: Callable[[Parameter], bool] | None = None,
diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index 4c79c63..da06d8e 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -5,16 +5,21 @@
 from pytensor import config
 
 from pytensor_ml.optim import (
+    adadelta,
+    adagrad,
     adam,
+    adamax,
+    adamw,
     chain,
     compile_train,
     cosine_schedule,
+    nadam,
+    rmsprop,
     rprop,
     rprop_updates,
     scalar_state,
-    scale_by_schedule,
+    scale,
     sgd,
-    substitute_schedule,
 )
 from pytensor_ml.params import StepCounter, step_counter, trainable
 from pytensor_ml.pytensorf import collect_step_counters
@@ -67,34 +72,32 @@ def test_schedule_rate_alone_sets_the_step_size():
     np.testing.assert_allclose(p.get_value(), [1.71], rtol=RTOL)
 
 
-def test_substitution_matches_terminal_scaling_for_a_multiplicative_rate():
-    """The two mechanisms must agree wherever the rate is a plain multiplier on the step, which holds for
-    every rule in the tree. Only their behaviour on rules that consume the rate elsewhere differs."""
+def test_a_scheduled_rate_matches_terminal_scaling_for_a_multiplicative_rate():
+    """A rate the rule reads and the same rate applied to the finished step must agree wherever the rate is a
+    plain multiplier, which holds for every rule in the tree."""
     schedule = cosine_schedule(0.05, 20)
 
-    substituted, substituted_loss = quadratic_problem()
-    substituted_step = compile_train(substituted_loss, adam(learning_rate=schedule))
+    read_by_the_rule, rule_loss = quadratic_problem()
+    rule_step = compile_train(rule_loss, adam(learning_rate=schedule))
 
     scaled, scaled_loss = quadratic_problem()
-    scaled_step = compile_train(
-        scaled_loss, chain(adam(learning_rate=1.0), scale_by_schedule(schedule))
-    )
+    clock = step_counter()
+    scaled_step = compile_train(scaled_loss, chain(adam(learning_rate=1.0), scale(schedule(clock))))
 
     for _ in range(10):
-        substituted_step()
+        rule_step()
         scaled_step()
 
-    np.testing.assert_allclose(substituted.get_value(), scaled.get_value(), rtol=1e-5)
+    np.testing.assert_allclose(read_by_the_rule.get_value(), scaled.get_value(), rtol=1e-5)
 
 
-def test_scheduled_rule_publishes_the_applied_rate():
+def test_a_scheduled_rule_allocates_no_rate_variable():
+    """The rate is an expression over the clock, so there is no rate state to keep or checkpoint. A caller
+    who wants to read it builds `curve(clock)` and asks for it as an extra output."""
     p, loss = quadratic_problem()
-    rule = adam(learning_rate=cosine_schedule(0.1, 4))
-    updates = rule(loss, [p])
-    published_rate = next(key for key in updates if key.name == "adam/learning_rate")
+    updates = adam(learning_rate=cosine_schedule(0.1, 4))(loss, [p])
 
-    compile_train(loss, rule)()
-    np.testing.assert_allclose(published_rate.get_value(), 0.1, rtol=RTOL)
+    assert not any("learning_rate" in (key.name or "") for key in updates)
 
 
 @pytest.mark.parametrize(
@@ -112,16 +115,7 @@ def test_rprop_rejects_a_symbolic_rate(learning_rate):
         rprop_updates(loss, [p], learning_rate=learning_rate)
 
 
-def test_substitute_schedule_raises_when_the_rate_is_absent():
-    p, loss = quadratic_problem()
-    absent_rate = scalar_state("absent/learning_rate", fill_value=0.1)
-    updates = sgd(learning_rate=0.1)(loss, [p])
-
-    with pytest.raises(ValueError, match="does not appear in the updates"):
-        substitute_schedule(updates, absent_rate, lambda step_count: pt.constant(0.05))
-
-
-def test_substitution_preserves_parameter_identity():
+def test_a_scheduled_rule_preserves_parameter_identity():
     """The updates must keep updating the parameter object the model holds; a clone would train a copy and
     silently leave the model's weights untouched."""
     p, loss = quadratic_problem()
@@ -134,8 +128,8 @@ def test_substitution_preserves_parameter_identity():
 
 
 def test_schedule_reaches_a_rate_applied_by_a_chained_scale():
-    """sgd with momentum applies its rate through `scale`, not through sgd_updates, so substitution has to
-    find it there. Unit-rate sgd gives step -p and momentum leaves the first step unchanged."""
+    """sgd with momentum applies its rate through `scale` rather than through sgd_updates, so the rate has to
+    reach it there. Unit-rate sgd gives step -p and momentum leaves the first step unchanged."""
     p, loss = quadratic_problem()
     step = compile_train(loss, sgd(learning_rate=cosine_schedule(0.1, 4), momentum=0.9))
 
@@ -153,22 +147,23 @@ def test_a_clock_read_by_a_schedule_drives_the_rate():
     assert int(clock.get_value()) == 1
 
 
-def test_a_clocked_schedule_matches_substitution_over_many_steps():
-    """The clock has to place the schedule at the same step the substituted counter would."""
+def test_a_caller_held_clock_matches_the_rule_own_clock_over_many_steps():
+    """Handing a rule an unapplied schedule and handing it the same schedule read off a clock of your own must
+    place the curve at the same step, or the two ways of scheduling drift apart."""
     schedule = cosine_schedule(0.05, 20)
 
-    substituted, substituted_loss = quadratic_problem()
-    substituted_step = compile_train(substituted_loss, adam(learning_rate=schedule))
+    rule_clocked, rule_loss = quadratic_problem()
+    rule_step = compile_train(rule_loss, adam(learning_rate=schedule))
 
-    clocked, clocked_loss = quadratic_problem()
+    caller_clocked, caller_loss = quadratic_problem()
     clock = step_counter()
-    clocked_step = compile_train(clocked_loss, adam(learning_rate=schedule(clock)))
+    caller_step = compile_train(caller_loss, adam(learning_rate=schedule(clock)))
 
     for _ in range(40):
-        substituted_step()
-        clocked_step()
+        rule_step()
+        caller_step()
 
-    np.testing.assert_allclose(clocked.get_value(), substituted.get_value(), rtol=RTOL)
+    np.testing.assert_allclose(caller_clocked.get_value(), rule_clocked.get_value(), rtol=RTOL)
 
 
 def test_a_clocked_rate_is_available_as_an_extra_output():
@@ -313,6 +308,16 @@ def test_a_rule_counts_its_own_steps_on_a_clock():
     assert all(isinstance(counter, StepCounter) for counter in counters)
 
 
+def test_a_scheduled_rule_keeps_one_clock():
+    """The schedule reads the clock the rule already counts on, so there is one notion of time to restore and
+    nothing that can disagree with itself."""
+    p, loss = quadratic_problem()
+    step = compile_train(loss, adam(learning_rate=cosine_schedule(0.1, 20)))
+
+    counters = [v for v in step.get_shared() if v.type.dtype.startswith("int")]
+    assert [counter.name for counter in counters] == ["adam/step_count"]
+
+
 def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
     """Two clocks can still coexist when the caller holds one of their own, and a checkpoint that restored
     only one of them leaves the two measuring different times."""
@@ -326,6 +331,25 @@ def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
         compile_train(loss, rule, extra_outputs=[their_clock.astype("float64")])
 
 
+@pytest.mark.parametrize(
+    "alias",
+    [adam, adamw, nadam, adamax, sgd, rmsprop, adagrad, adadelta],
+    ids=["adam", "adamw", "nadam", "adamax", "sgd", "rmsprop", "adagrad", "adadelta"],
+)
+def test_a_scheduled_rule_keeps_one_notion_of_time(alias):
+    """A rule that counts its own steps and also reads a schedule can end up holding more than one clock,
+    which is harmless as long as they agree: each advances once per step from zero, so the rate the schedule
+    computes and the bias correction the rule applies are always at the same step."""
+    p, loss = quadratic_problem()
+    step = compile_train(loss, alias(learning_rate=cosine_schedule(0.1, 20)))
+
+    for _ in range(3):
+        step()
+
+    clocks = [v for v in step.get_shared() if isinstance(v, StepCounter)]
+    assert clocks and {int(clock.get_value()) for clock in clocks} == {3}
+
+
 def test_a_caller_held_clock_keeps_advancing_for_other_readers():
     """A clock is state that outlives the compile: the caller reads it, and so can another compiled function.
     Compiling a training step must not take it out of the set that step advances."""
diff --git a/tests/optim/test_schedules.py b/tests/optim/test_schedules.py
index 93dacc9..3f7ab37 100644
--- a/tests/optim/test_schedules.py
+++ b/tests/optim/test_schedules.py
@@ -13,11 +13,11 @@
     join_schedules,
     linear_schedule,
     polynomial_schedule,
-    scale_by_schedule,
+    scale,
     sgd,
     step_decay,
 )
-from pytensor_ml.params import trainable
+from pytensor_ml.params import step_counter, trainable
 from pytensor_ml.pytensorf import function
 
 # Every schedule takes (learning_rate, total_steps, final_learning_rate) and owes the same contract at the
@@ -178,26 +178,25 @@ def test_schedule_rejects_negative_transition_begin(schedule_factory):
 
 @pytest.mark.parametrize("schedule_factory", SCHEDULES)
 def test_schedule_drives_training_through_the_learning_rate_union(schedule_factory):
-    """Passing a schedule as `learning_rate` substitutes it into the rule's own rate, which is a different
-    path from `scale_by_schedule` and the one a new schedule is most likely to miss."""
+    """Passing a schedule as `learning_rate` reads it off the clock the rule counts on, which is the path a
+    new schedule is most likely to miss."""
     p = trainable(np.array([2.0]), name="w")
     loss = 0.5 * (p**2).sum()  # grad = p, so the unit-rate base step is -p
     rule = sgd(learning_rate=schedule_factory(0.1, 2, 0.01))
-    published_rate = next(key for key in rule(loss, [p]) if key.name == "sgd/learning_rate")
 
     compile_train(loss, rule)()  # every schedule starts at its initial rate, so p = 2 - 0.1 * 2
-    np.testing.assert_allclose(published_rate.get_value(), 0.1, rtol=1e-6)
     np.testing.assert_allclose(p.get_value(), [1.8], rtol=1e-6)
 
 
 # Geometric decay never reaches zero, so the schedule that cannot hit a zero floor sits this one out.
 @pytest.mark.parametrize("schedule_factory", [cosine_schedule, linear_schedule])
-def test_schedule_drives_training_through_scale_by_schedule(schedule_factory):
+def test_schedule_drives_training_through_a_scaled_chain(schedule_factory):
     # Decaying to zero over two steps, every curve passes through the same midpoint, so one set of expected
     # values covers them: 0.1 -> 0.05 -> 0.
     p = trainable(np.array([2.0]), name="w")
     loss = 0.5 * (p**2).sum()  # grad = p, so the unit-rate base step is -p
-    rule = chain(sgd(learning_rate=1.0), scale_by_schedule(schedule_factory(0.1, 2)))
+    clock = step_counter()
+    rule = chain(sgd(learning_rate=1.0), scale(schedule_factory(0.1, 2)(clock)))
     step = compile_train(loss, rule)
 
     step()  # step 0 -> lr = 0.1, p = 2 - 0.1 * 2 = 1.8
diff --git a/tests/optim/test_transform.py b/tests/optim/test_transform.py
index 5b897b8..376c9ae 100644
--- a/tests/optim/test_transform.py
+++ b/tests/optim/test_transform.py
@@ -1,18 +1,7 @@
 import numpy as np
 import pytensor.tensor as pt
-import pytest
-
-from pytensor import config
-
-from pytensor_ml.optim import (
-    add_weight_decay,
-    chain,
-    scalar_state,
-    scale,
-    scale_by_schedule,
-    sgd_updates,
-    trace,
-)
+
+from pytensor_ml.optim import add_weight_decay, chain, scale, sgd_updates, trace
 from pytensor_ml.params import trainable
 from pytensor_ml.pytensorf import function
 
@@ -61,28 +50,12 @@ def test_chain_threads_updates_in_order():
     np.testing.assert_allclose(function([], out[p])(), [0.9])
 
 
-def test_scale_by_schedule_applies_decaying_rate():
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()  # grad = p, so the unit-rate base step is -p
-
-    def schedule(step_count):
-        return 0.1 / (1.0 + step_count.astype("float64"))
-
-    out = scale_by_schedule(schedule)(sgd_updates(loss, [p], learning_rate=1.0), [p])
-    step = function([], loss, updates=out)
-
-    step()  # step_count=0 -> lr=0.1, step=-2, p = 2 - 0.1 * 2 = 1.8
-    np.testing.assert_allclose(p.get_value(), [1.8], rtol=1e-6)
-    step()  # step_count=1 -> lr=0.05, step=-1.8, p = 1.8 - 0.05 * 1.8 = 1.71
-    np.testing.assert_allclose(p.get_value(), [1.71], rtol=1e-6)
-
-
 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 the schedule with its own step counter."""
+    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_by_schedule(lambda step_count: pt.constant(0.1)))
+    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}
@@ -95,91 +68,12 @@ def test_separately_configured_chains_keep_independent_state():
     loss = 0.5 * (p**2).sum()
 
     def build_updates():
-        rule = chain(trace(0.9), scale_by_schedule(lambda step_count: pt.constant(0.1)))
+        rule = chain(trace(0.9), scale(0.1))
         return {key for key in rule(sgd_updates(loss, [p], learning_rate=1.0), [p]) if key is not p}
 
     assert not build_updates() & build_updates()
 
 
-def test_scale_by_schedule_publishes_the_applied_rate():
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()
-    out = scale_by_schedule(lambda step_count: 0.1 / (1.0 + step_count.astype(config.floatX)))(
-        sgd_updates(loss, [p], learning_rate=1.0), [p]
-    )
-    published_rate = next(key for key in out if key.name == "schedule/learning_rate")
-    step = function([], loss, updates=out)
-
-    # The rate published is the rate applied this step, not the one used on the step before: p goes
-    # 2 -> 1.8 under lr=0.1, and the variable holds 0.1 rather than its initial 0.
-    step()
-    np.testing.assert_allclose(published_rate.get_value(), 0.1, rtol=1e-6)
-    np.testing.assert_allclose(p.get_value(), [1.8], rtol=1e-6)
-    step()
-    np.testing.assert_allclose(published_rate.get_value(), 0.05, rtol=1e-6)
-
-
-def test_scale_by_schedule_publishes_to_a_caller_held_variable():
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()
-    learning_rate = scalar_state("my/learning_rate")
-    out = scale_by_schedule(lambda step_count: pt.constant(0.25), learning_rate=learning_rate)(
-        sgd_updates(loss, [p], learning_rate=1.0), [p]
-    )
-    step = function([], loss, updates=out)
-
-    step()
-    np.testing.assert_allclose(learning_rate.get_value(), 0.25, rtol=1e-6)
-    assert not any(key.name == "schedule/learning_rate" for key in out)
-
-
-def test_two_scheduled_scalings_in_one_chain_raise():
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()
-    constant_scaling = scale_by_schedule(lambda step_count: pt.constant(0.1))
-    with pytest.raises(ValueError, match="Two scheduled scalings in one chain"):
-        chain(constant_scaling, constant_scaling)(sgd_updates(loss, [p], learning_rate=1.0), [p])
-
-
-def test_scale_by_schedule_casts_rate_to_parameter_dtype():
-    # A float64 schedule must not upcast a float32 parameter's update, which pytensor rejects outright.
-    with config.change_flags(floatX="float32"):
-        p = trainable(np.array([2.0], dtype="float32"), name="w")
-        loss = 0.5 * (p**2).sum()
-        out = scale_by_schedule(lambda step_count: 0.1 / (1.0 + step_count.astype("float64")))(
-            sgd_updates(loss, [p], learning_rate=1.0), [p]
-        )
-        published_rate = next(key for key in out if key.name == "schedule/learning_rate")
-        assert out[p].type.dtype == "float32"
-        assert out[published_rate].type.dtype == "float32"
-
-        step = function([], loss, updates=out)
-        step()
-        np.testing.assert_allclose(p.get_value(), [1.8], rtol=1e-6)
-
-
-def test_two_scheduled_scalings_compose_when_given_distinct_variables():
-    """The workaround the rejection above recommends: distinct rate variables let two schedules stack.
-    Each publishes its own factor, they share one step counter, and the parameters see the product."""
-    p = trainable(np.array([2.0]), name="w")
-    loss = 0.5 * (p**2).sum()
-    warmup_rate = scalar_state("warmup/learning_rate")
-    decay_rate = scalar_state("decay/learning_rate")
-
-    out = chain(
-        scale_by_schedule(lambda step_count: pt.constant(0.5), learning_rate=warmup_rate),
-        scale_by_schedule(lambda step_count: pt.constant(0.2), learning_rate=decay_rate),
-    )(sgd_updates(loss, [p], learning_rate=1.0), [p])
-    step_count = next(key for key in out if key.name == "schedule/step_count")
-    step = function([], loss, updates=out)
-
-    step()  # the applied rate is 0.5 * 0.2, so p = 2 - 0.1 * 2
-    np.testing.assert_allclose(p.get_value(), [1.8], rtol=1e-6)
-    np.testing.assert_allclose(warmup_rate.get_value(), 0.5, rtol=1e-6)
-    np.testing.assert_allclose(decay_rate.get_value(), 0.2, rtol=1e-6)
-    assert int(step_count.get_value()) == 1  # one counter, advanced once: both schedules see one t
-
-
 def test_add_weight_decay_subtracts_decay_term():
     p = trainable(np.array([4.0]), name="w")
     # An empty base step (updates[p] == p) isolates the decay term: new step = -0.1 * 4.

From da6591998de50fd8696e26c278e76f38fe4ff7fa Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Sun, 9 Aug 2026 19:28:23 -0500
Subject: [PATCH 3/5] Assert a scheduled rule keeps exactly one clock

Every rule already counted its steps under the name its alias asks for, so the two-clock case the weaker assertion allowed for cannot arise.
---
 pytensor_ml/optim/alias.py        |  6 +++---
 tests/optim/test_learning_rate.py | 36 +++++++++++++++----------------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py
index 6f70e4d..2cf9229 100644
--- a/pytensor_ml/optim/alias.py
+++ b/pytensor_ml/optim/alias.py
@@ -30,9 +30,9 @@ def _at_learning_rate(
     name: str,
     build_updates: Callable[[Rate], Updates],
 ) -> Updates:
-    """Build updates at ``learning_rate``, reading a schedule off the rule's step clock. A rule that keeps its
-    own clock under another name ends up with two, which is harmless: every clock advances once per step from
-    zero, so they agree about the step at every point."""
+    """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)
 
diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index da06d8e..77b771b 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -308,16 +308,6 @@ def test_a_rule_counts_its_own_steps_on_a_clock():
     assert all(isinstance(counter, StepCounter) for counter in counters)
 
 
-def test_a_scheduled_rule_keeps_one_clock():
-    """The schedule reads the clock the rule already counts on, so there is one notion of time to restore and
-    nothing that can disagree with itself."""
-    p, loss = quadratic_problem()
-    step = compile_train(loss, adam(learning_rate=cosine_schedule(0.1, 20)))
-
-    counters = [v for v in step.get_shared() if v.type.dtype.startswith("int")]
-    assert [counter.name for counter in counters] == ["adam/step_count"]
-
-
 def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
     """Two clocks can still coexist when the caller holds one of their own, and a checkpoint that restored
     only one of them leaves the two measuring different times."""
@@ -332,22 +322,32 @@ def test_a_user_clock_out_of_step_with_a_rule_clock_is_caught():
 
 
 @pytest.mark.parametrize(
-    "alias",
-    [adam, adamw, nadam, adamax, sgd, rmsprop, adagrad, adadelta],
+    "alias, clock_name",
+    [
+        (adam, "adam/step_count"),
+        (adamw, "adamw/step_count"),
+        (nadam, "nadam/step_count"),
+        (adamax, "adamax/step_count"),
+        (sgd, "sgd/step_count"),
+        (rmsprop, "rmsprop/step_count"),
+        (adagrad, "adagrad/step_count"),
+        (adadelta, "adadelta/step_count"),
+    ],
     ids=["adam", "adamw", "nadam", "adamax", "sgd", "rmsprop", "adagrad", "adadelta"],
 )
-def test_a_scheduled_rule_keeps_one_notion_of_time(alias):
-    """A rule that counts its own steps and also reads a schedule can end up holding more than one clock,
-    which is harmless as long as they agree: each advances once per step from zero, so the rate the schedule
-    computes and the bias correction the rule applies are always at the same step."""
-    p, loss = quadratic_problem()
+def test_a_scheduled_rule_keeps_exactly_one_clock(alias, clock_name):
+    """One notion of time per rule, under the name the schedule and the rule's own step count agree on. A rule
+    that counted its steps under a different name would hold a second clock measuring the same time, and a
+    checkpoint would carry two counts to keep in step."""
+    _, loss = quadratic_problem()
     step = compile_train(loss, alias(learning_rate=cosine_schedule(0.1, 20)))
 
     for _ in range(3):
         step()
 
     clocks = [v for v in step.get_shared() if isinstance(v, StepCounter)]
-    assert clocks and {int(clock.get_value()) for clock in clocks} == {3}
+    assert [clock.name for clock in clocks] == [clock_name]
+    assert int(clocks[0].get_value()) == 3
 
 
 def test_a_caller_held_clock_keeps_advancing_for_other_readers():

From 81743582a5877f82abc40ce304855e3e08067d79 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Sun, 9 Aug 2026 19:29:00 -0500
Subject: [PATCH 4/5] Check the exact rate a second clock reader sees

The old bound passed with the clock at 1, where the rate is 0.9998, so it did not distinguish an advancing clock from a frozen one.
---
 tests/optim/test_learning_rate.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/tests/optim/test_learning_rate.py b/tests/optim/test_learning_rate.py
index 77b771b..5800713 100644
--- a/tests/optim/test_learning_rate.py
+++ b/tests/optim/test_learning_rate.py
@@ -22,8 +22,7 @@
     sgd,
 )
 from pytensor_ml.params import StepCounter, step_counter, trainable
-from pytensor_ml.pytensorf import collect_step_counters
-from pytensor_ml.pytensorf import function as pytensor_function
+from pytensor_ml.pytensorf import collect_step_counters, function
 
 RTOL = 1e-6
 
@@ -353,13 +352,15 @@ def test_a_scheduled_rule_keeps_exactly_one_clock(alias, clock_name):
 def test_a_caller_held_clock_keeps_advancing_for_other_readers():
     """A clock is state that outlives the compile: the caller reads it, and so can another compiled function.
     Compiling a training step must not take it out of the set that step advances."""
-    p, loss = quadratic_problem()
+    _, loss = quadratic_problem()
     clock = step_counter("agent_clock")
     step = compile_train(loss, adam(learning_rate=cosine_schedule(0.1, 100)(clock)))
-    epsilon = pytensor_function([], cosine_schedule(1.0, 100)(clock))
+    exploration_schedule = cosine_schedule(1.0, 100)
+    epsilon = function([], exploration_schedule(clock))
 
     for _ in range(20):
         step()
 
     assert int(clock.get_value()) == 20
-    assert float(epsilon()) < 1.0
+    expected = float(exploration_schedule(pt.as_tensor(20, dtype="int64")).eval())
+    np.testing.assert_allclose(epsilon(), expected, rtol=RTOL)

From 010c8e169f891fece42fcf562786f3baad1d9b78 Mon Sep 17 00:00:00 2001
From: jessegrabowski <jessegrabowski@gmail.com>
Date: Sun, 9 Aug 2026 19:29:16 -0500
Subject: [PATCH 5/5] Trim the repeated schedule return description

---
 pytensor_ml/optim/schedules.py | 35 ++++++++++++++--------------------
 1 file changed, 14 insertions(+), 21 deletions(-)

diff --git a/pytensor_ml/optim/schedules.py b/pytensor_ml/optim/schedules.py
index 6b424cd..182cddf 100644
--- a/pytensor_ml/optim/schedules.py
+++ b/pytensor_ml/optim/schedules.py
@@ -61,9 +61,8 @@ def cosine_schedule(
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
 
     Examples
     --------
@@ -126,9 +125,8 @@ def linear_schedule(
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
     _validate_horizon(total_steps, transition_begin)
 
@@ -181,9 +179,8 @@ def exponential_schedule(
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
     _validate_horizon(total_steps, transition_begin)
     if final_learning_rate <= 0.0 or learning_rate <= 0.0:
@@ -245,9 +242,8 @@ def polynomial_schedule(
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
     _validate_horizon(total_steps, transition_begin)
     if power <= 0.0:
@@ -304,9 +300,8 @@ def step_decay(
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
     if decay_every < 1:
         raise ValueError(f"decay_every must be at least 1, got {decay_every}.")
@@ -342,9 +337,8 @@ def constant_schedule(learning_rate: float) -> Schedule:
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
 
     def schedule(step_count: TensorVariable) -> TensorVariable:
@@ -380,9 +374,8 @@ def join_schedules(schedules: Sequence[Schedule], boundaries: Sequence[int]) ->
     Returns
     -------
     Schedule
-        A callable mapping a symbolic step count to a scalar learning rate. Apply it to a
-        :func:`~pytensor_ml.params.step_counter` to get the rate as a graph, or hand it to a rule as its
-        ``learning_rate`` to have the rule read it off the clock it already counts on.
+        A callable mapping a symbolic step count to a scalar learning rate, ready to hand to a rule as
+        its ``learning_rate``.
     """
     if not schedules:
         raise ValueError("join_schedules needs at least one schedule.")