Memoize optimizer state per rule and parameter

pytensor_ml/optim/alias.py

@@ -1,6 +1,12 @@
 from collections.abc import Callable, Sequence
 
-from pytensor_ml.optim.base import LossOrGradients, Parameter, UpdateRule, Updates
+from pytensor_ml.optim.base import (
+    LossOrGradients,
+    Parameter,
+    UpdateRule,
+    Updates,
+    reuses_state,
+)
 from pytensor_ml.optim.rules import (
     adadelta_updates,
     adagrad_updates,
@@ -29,6 +35,7 @@ def sgd(learning_rate: float = 0.01, momentum: float = 0.0, nesterov: bool = Fal
         Use Nesterov momentum. Ignored when ``momentum`` is 0. Default False.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         if not momentum:
             return sgd_updates(loss_or_gradients, parameters, learning_rate=learning_rate)
@@ -50,6 +57,7 @@ def adam(
     Adam optimizer. See :func:`~pytensor_ml.optim.rules.adam_updates` for the update rule.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return adam_updates(
             loss_or_gradients,
@@ -78,6 +86,7 @@ def adamw(
     :func:`~pytensor_ml.optim.rules.adamw_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return adamw_updates(
             loss_or_gradients,
@@ -105,6 +114,7 @@ def nadam(
     :func:`~pytensor_ml.optim.rules.nadam_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return nadam_updates(
             loss_or_gradients,
@@ -129,6 +139,7 @@ def adamax(
     :func:`~pytensor_ml.optim.rules.adamax_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return adamax_updates(
             loss_or_gradients,
@@ -153,6 +164,7 @@ def rprop(
     Rprop optimizer (resilient backpropagation). See :func:`~pytensor_ml.optim.rules.rprop_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return rprop_updates(
             loss_or_gradients,
@@ -178,6 +190,7 @@ def rmsprop(
     RMSProp optimizer. See :func:`~pytensor_ml.optim.rules.rmsprop_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return rmsprop_updates(
             loss_or_gradients,
@@ -197,6 +210,7 @@ def adagrad(learning_rate: float = 0.01, epsilon: float = 1e-8) -> UpdateRule:
     AdaGrad optimizer. See :func:`~pytensor_ml.optim.rules.adagrad_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return adagrad_updates(
             loss_or_gradients, parameters, learning_rate=learning_rate, epsilon=epsilon
@@ -210,6 +224,7 @@ def adadelta(learning_rate: float = 1.0, rho: float = 0.9, epsilon: float = 1e-8
     AdaDelta optimizer. See :func:`~pytensor_ml.optim.rules.adadelta_updates`.
     """
 
+    @reuses_state
     def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates:
         return adadelta_updates(
             loss_or_gradients, parameters, learning_rate=learning_rate, rho=rho, epsilon=epsilon

pytensor_ml/optim/base.py

@@ -1,4 +1,6 @@
 from collections.abc import Callable, Sequence
+from contextvars import ContextVar
+from functools import wraps
 
 import numpy as np
 import pytensor
@@ -52,13 +54,67 @@ def get_gradients(
     return grad(rewrite_pregrad(loss_or_gradients), list(parameters))  # type: ignore[return-value]
 
 
+# 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
+)
+
+
+def reuses_state(rule: UpdateRule) -> UpdateRule:
+    """
+    Give ``rule`` 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 buffers are keyed per rule rather than globally so two independently configured optimizers stay
+    independent. They are bound dynamically because :func:`state_for` is reached through the update-rule
+    functions, several call layers below, and threading a cache down would touch every one of them.
+
+    Parameters
+    ----------
+    rule : UpdateRule
+        The rule to wrap. Its buffers live as long as the wrapper does.
+    """
+    buffers: dict[_StateKey, Parameter] = {}
+
+    @wraps(rule)
+    def rule_with_persistent_state(
+        loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]
+    ) -> Updates:
+        token = _state_buffers.set(buffers)
+        try:
+            return rule(loss_or_gradients, parameters)
+        finally:
+            _state_buffers.reset(token)
+
+    return rule_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:
     """
-    Allocate an optimizer-state shared variable shaped and typed like ``parameter``.
+    Return the optimizer-state shared variable shaped and typed like ``parameter``.
 
     The variable is named ``"{parameter.name}/{slot}"`` so it can be matched by name at serialization
     boundaries. The name is never used to *find* the variable at runtime — callers hold the returned object
-    directly.
+    directly, and reuse within a rule is keyed on the parameter object, so two same-named parameters still
+    get distinct buffers and collide loudly at save time rather than silently sharing.
+
+    Allocates unless the enclosing rule was wrapped in :func:`reuses_state` and already holds this slot.
 
     Parameters
     ----------
@@ -68,19 +124,26 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param
         A short role tag for the slot, e.g. ``"adam/first_moment"`` or ``"trace/velocity"``.
     fill_value : float
         Constant to initialize the state with. Default 0.0.
-
-    Returns
-    -------
-    shared tensor variable
-        A freshly allocated state variable.
     """
     if parameter.name is None:
         raise ValueError(
             f"Cannot allocate optimizer state {slot!r} for an unnamed parameter. Stateful optimizers rely on "
             "parameter names to identify their state at serialization boundaries; give the parameter a name."
         )
-    value = parameter.get_value(borrow=True)
-    return pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}")
+
+    def allocate() -> Parameter:
+        value = parameter.get_value(borrow=True)
+        return pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}")
+
+    return _reuse_or_allocate((parameter, slot), allocate)
+
+
+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)
+    )
 
 
 def require_unique_state_names(updates: Updates) -> None:
@@ -109,11 +172,6 @@ def require_unique_state_names(updates: Updates) -> None:
         seen.add(name)
 
 
-def counter(name: str) -> Parameter:
-    """Allocate an int64 step counter shared variable initialized to zero."""
-    return pytensor.shared(np.asarray(0, dtype="int64"), name=name)
-
-
 def chain(*transforms: Transform) -> Transform:
     """
     Compose transforms left to right into a single transform.

tests/optim/test_rules.py

@@ -152,6 +152,59 @@ def test_adam_updates_keyed_by_object_with_named_state():
     assert state_names == {"adam/step_count", "w/adam/first_moment", "w/adam/second_moment"}
 
 
+@pytest.mark.parametrize(
+    "make_rule",
+    [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."""
+    p = trainable(np.zeros(3), name="w")
+    loss = (p**2).sum()
+    rule = 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
+
+
+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``."""
+    p = trainable(np.zeros(2), name="w")
+    gradient = np.array([2.0, -0.5])
+    loss = (pt.constant(gradient) * p).sum()  # constant gradient, independent of p
+    learning_rate, momentum = 0.1, 0.9
+    rule = sgd(learning_rate=learning_rate, momentum=momentum)
+
+    step_once = function([], loss, updates=rule(loss, [p]))
+    step_again = function([], loss, updates=rule(loss, [p]))
+
+    step_once()
+    before = p.get_value().copy()
+    step_again()
+
+    continued = -learning_rate * gradient * (1 - momentum**2) / (1 - momentum)
+    np.testing.assert_allclose(p.get_value() - before, continued, rtol=1e-6)
+
+
+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