fix(alias): un-break downstream testing for weight_decay default 0.0 param requirements

Author: rdyro

optax/transforms/_adding.py

@@ -57,8 +57,6 @@ def init_fn(params):
       return base.EmptyState()
 
   def update_fn(updates, state, params):
-    if params is None:
-      raise ValueError(base.NO_PARAMS_MSG)
     if callable(weight_decay):
       new_state = WeightDecaySchedule(numerics.safe_increment(state.count))
     else:
@@ -68,6 +66,9 @@ def update_fn(updates, state, params):
     if isinstance(weight_decay, (int, float)) and weight_decay == 0.0:
       return updates, new_state
 
+    if params is None:
+      raise ValueError(base.NO_PARAMS_MSG)
+
     s = weight_decay(state.count) if callable(weight_decay) else weight_decay
     updates = jax.tree.map(
         lambda g, p: None if g is None else g + s * p,

optax/tree_utils/_state_utils_test.py

@@ -195,6 +195,7 @@ def test_map_params_to_none(self):
         (
             transform.ScaleByRssState(sum_of_squares={'a': None}),
             base.EmptyState(),
+            base.EmptyState(),
         ),
     )
 
@@ -378,7 +379,7 @@ def get_learning_rate(state):
 
       for i in range(4):
         # we simply update state, we don't care about updates.
-        _, state = opt.update(params, state)
+        _, state = opt.update(params, state, params)
         lr = get_learning_rate(state)
         self.assertEqual(lr, 1 / (i + 1))
 
@@ -465,7 +466,7 @@ def set_learning_rate(state, lr):
       for i in range(4):
         modified_state = set_learning_rate(modified_state, lr / (i + 1))
         # we simply update state, we don't care about updates.
-        _, modified_state = opt.update(params, modified_state)
+        _, modified_state = opt.update(params, modified_state, params)
         modified_lr = _state_utils.tree_get(modified_state, 'learning_rate')
         self.assertEqual(modified_lr, lr / (i + 1))