Allocate optimizer state afresh on every rule invocation

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

docs/source/api/optim.rst

@@ -86,9 +86,7 @@ Building blocks
 .. autosummary::
     :toctree: generated/
 
-    counter
     get_gradients
-    reuses_state
     scalar_state
     state_for
     steps_of

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

pytensor_ml/optim/__init__.py

@@ -18,9 +18,7 @@
     Transform,
     Updates,
     chain,
-    counter,
     get_gradients,
-    reuses_state,
     scalar_state,
     state_for,
     steps_of,
@@ -94,7 +92,6 @@
     "compile_train",
     "constant_schedule",
     "cosine_schedule",
-    "counter",
     "exponential_schedule",
     "get_gradients",
     "join_schedules",
@@ -106,7 +103,6 @@
     "nonfinite",
     "polynomial_schedule",
     "reduce_on_plateau",
-    "reuses_state",
     "rmsprop",
     "rmsprop_updates",
     "rprop",

pytensor_ml/optim/alias.py

@@ -4,11 +4,8 @@
     LearningRate,
     LossGradientsOrUpdates,
     Parameter,
-    Rate,
     Transform,
     Updates,
-    counter,
-    reuses_state,
 )
 from pytensor_ml.optim.rules import (
     _require_numeric_learning_rate,
@@ -25,20 +22,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,
@@ -86,32 +69,24 @@ 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:
-        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 not momentum:
+            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 = trace(momentum, nesterov, namespace=namespace)(updates, parameters)
+        return scale(learning_rate, namespace=namespace)(updates, parameters)
 
     return rule
 
@@ -151,23 +126,18 @@ 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:
-        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
@@ -211,25 +181,20 @@ 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:
-        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
@@ -270,22 +235,17 @@ 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:
-        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
@@ -326,22 +286,17 @@ 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:
-        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
@@ -385,7 +340,6 @@ def rprop(
     """
     _require_numeric_learning_rate(learning_rate)
 
-    @reuses_state
     def rule(
         loss_gradients_or_updates: LossGradientsOrUpdates, parameters: Sequence[Parameter]
     ) -> Updates:
@@ -438,23 +392,18 @@ 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:
-        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
@@ -492,20 +441,15 @@ 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:
-        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
@@ -544,21 +488,16 @@ 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:
-        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