Read schedules off the clock the rule counts on

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",
 ]

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,21 +22,21 @@
     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(
     learning_rate: LearningRate,
     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 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)
 
-    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

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:

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

pytensor_ml/optim/schedules.py

@@ -61,8 +61,8 @@ 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, ready to hand to a rule as
+        its ``learning_rate``.
 
     Examples
     --------
@@ -125,8 +125,8 @@ 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, ready to hand to a rule as
+        its ``learning_rate``.
     """
     _validate_horizon(total_steps, transition_begin)
 
@@ -179,8 +179,8 @@ 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, 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:
@@ -242,8 +242,8 @@ 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, ready to hand to a rule as
+        its ``learning_rate``.
     """
     _validate_horizon(total_steps, transition_begin)
     if power <= 0.0:
@@ -300,8 +300,8 @@ 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, 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}.")
@@ -337,8 +337,8 @@ 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, ready to hand to a rule as
+        its ``learning_rate``.
     """
 
     def schedule(step_count: TensorVariable) -> TensorVariable:
@@ -374,8 +374,8 @@ 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, ready to hand to a rule as
+        its ``learning_rate``.
     """
     if not schedules:
         raise ValueError("join_schedules needs at least one schedule.")

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)
 

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,

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}