Ensure optimizers accept extra args

Most already use optax.chain which enforces this
Fix for #1131

PiperOrigin-RevId: 731893013

Author: rdyro

optax/_src/alias.py

@@ -39,7 +39,7 @@ def adabelief(
     eps_root: float = 1e-16,
     *,
     nesterov: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The AdaBelief optimizer.
 
   AdaBelief is an adaptive learning rate optimizer that focuses on fast
@@ -141,7 +141,7 @@ def adadelta(
     eps: float = 1e-6,
     weight_decay: float = 0.0,
     weight_decay_mask: MaskOrFn = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The Adadelta optimizer.
 
   Adadelta is a stochastic gradient descent method that adapts learning rates
@@ -208,7 +208,7 @@ def adafactor(
     eps: float = 1e-30,
     factored: bool = True,
     weight_decay_mask: MaskOrFn = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The Adafactor optimizer.
 
   Adafactor is an adaptive learning rate optimizer that focuses on fast
@@ -304,7 +304,7 @@ def adagrad(
     learning_rate: base.ScalarOrSchedule,
     initial_accumulator_value: float = 0.1,
     eps: float = 1e-7,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The Adagrad optimizer.
 
   AdaGrad is a sub-gradient algorithm for stochastic optimization that adapts
@@ -394,7 +394,7 @@ def adam(
     mu_dtype: Optional[Any] = None,
     *,
     nesterov: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The Adam optimizer.
 
   Adam is an SGD variant with gradient scaling adaptation. The scaling
@@ -580,7 +580,7 @@ def adamw(
     mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
     *,
     nesterov: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""Adam with weight decay regularization.
 
   AdamW uses weight decay to regularize learning towards small weights, as
@@ -789,7 +789,7 @@ def adan(
     eps_root: float = 1e-8,
     weight_decay: float = 0.0,
     mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The ADAptive Nesterov momentum algorithm (Adan).
 
   Adan first reformulates the vanilla Nesterov acceleration to develop a new
@@ -905,7 +905,7 @@ def lion(
     mu_dtype: Optional[Any] = None,
     weight_decay: float = 1e-3,
     mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The Lion optimizer.
 
   Lion is discovered by symbolic program search. Unlike most adaptive optimizers
@@ -1001,7 +1001,7 @@ def amsgrad(
     eps: float = 1e-8,
     eps_root: float = 0.0,
     mu_dtype: Optional[Any] = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The AMSGrad optimizer.
 
   The original Adam can fail to converge to the optimal solution in some cases.
@@ -1058,7 +1058,7 @@ def amsgrad(
 
 def fromage(
     learning_rate: float, min_norm: float = 1e-6
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The Frobenius matched gradient descent (Fromage) optimizer.
 
   Fromage is a learning algorithm that does not require learning rate tuning.
@@ -1119,7 +1119,7 @@ def lars(
     trust_ratio_mask: MaskOrFn = True,
     momentum: float = 0.9,
     nesterov: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The LARS optimizer.
 
   LARS is a layer-wise adaptive optimizer introduced to help scale SGD to
@@ -1191,7 +1191,7 @@ def lamb(
     eps_root: float = 0.0,
     weight_decay: float = 0.0,
     mask: MaskOrFn = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The LAMB optimizer.
 
   LAMB is a general purpose layer-wise adaptive large batch optimizer designed
@@ -1257,7 +1257,7 @@ def noisy_sgd(
     eta: float = 0.01,
     gamma: float = 0.55,
     seed: int = 0,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""A variant of SGD with added noise.
 
   Noisy SGD is a variant of :func:`optax.sgd` that incorporates Gaussian noise
@@ -1325,7 +1325,7 @@ def noisy_sgd(
 
 def sign_sgd(
     learning_rate: base.ScalarOrSchedule,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""A variant of SGD using only the signs of the gradient components.
 
   SignSGD is a variant of SGD that uses the signs of the gradient components in
@@ -1394,7 +1394,7 @@ def novograd(
     eps: float = 1e-6,
     eps_root: float = 0.0,
     weight_decay: float = 0.0,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """NovoGrad optimizer.
 
   NovoGrad is more robust to the initial learning rate and
@@ -1461,7 +1461,7 @@ def optimistic_gradient_descent(
     learning_rate: base.ScalarOrSchedule,
     alpha: base.ScalarOrSchedule = 1.0,
     beta: base.ScalarOrSchedule = 1.0,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """An Optimistic Gradient Descent optimizer.
 
   Optimistic gradient descent is an approximation of extra-gradient methods
@@ -1523,7 +1523,7 @@ def optimistic_adam(
     mu_dtype: Optional[Any] = None,
     *,
     nesterov: bool = True,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The Optimistic Adam optimizer.
 
   This is an optimistic version of the Adam optimizer. It addresses the issue
@@ -1643,7 +1643,7 @@ def radam(
     threshold: float = 5.0,
     *,
     nesterov: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The Rectified Adam optimizer.
 
   The adaptive learning rate in Adam has undesirably large variance in early
@@ -1715,7 +1715,7 @@ def rmsprop(
     momentum: Optional[float] = None,
     nesterov: bool = False,
     bias_correction: bool = False,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""A flexible RMSProp optimizer.
 
   RMSProp is an SGD variant with learning rate adaptation. The `learning_rate`
@@ -1824,7 +1824,7 @@ def sgd(
     momentum: Optional[float] = None,
     nesterov: bool = False,
     accumulator_dtype: Optional[Any] = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""A canonical Stochastic Gradient Descent optimizer.
 
   This implements stochastic gradient descent. It also includes support for
@@ -1911,7 +1911,7 @@ def sgd(
 
 def sm3(
     learning_rate: float, momentum: float = 0.9
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""The SM3 optimizer.
 
   SM3 (Square-root of Minima of Sums of Maxima of Squared-gradients Method) is a
@@ -2024,7 +2024,7 @@ def yogi(
     b1: float = 0.9,
     b2: float = 0.999,
     eps: float = 1e-3,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   # pylint: disable=line-too-long
   """The Yogi optimizer.
 
@@ -2083,7 +2083,7 @@ def adamax(
     b1: float = 0.9,
     b2: float = 0.999,
     eps: float = 1e-8,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   r"""A variant of the Adam optimizer that uses the infinity norm.
 
   AdaMax is a variant of the :func:`optax.adam` optimizer. By generalizing
@@ -2170,7 +2170,7 @@ def adamaxw(
     eps: float = 1e-8,
     weight_decay: float = 1e-4,
     mask: Optional[Union[Any, Callable[[base.Params], Any]]] = None,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """Adamax with weight decay regularization.
 
   AdamaxW uses weight decay to regularize learning towards small weights, as
@@ -2244,7 +2244,7 @@ def rprop(
     eta_plus: float = 1.2,
     min_step_size: float = 1e-6,
     max_step_size: float = 50.0,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """The Rprop optimizer.
 
   Rprop, short for resillient backpropogation, is a first order variant of
@@ -2405,7 +2405,7 @@ def lbfgs(
     memory_size: int = 10,
     scale_init_precond: bool = True,
     linesearch: Optional[
-        base.GradientTransformationExtraArgs
+        Union[base.GradientTransformationExtraArgs, base.GradientTransformation]
     ] = _linesearch.scale_by_zoom_linesearch(
         max_linesearch_steps=20, initial_guess_strategy='one'
     ),

optax/_src/alias_test.py

@@ -187,6 +187,30 @@ def step(params, state):
 
       chex.assert_trees_all_close(params, final_params, rtol=3e-2, atol=3e-2)
 
+  @parameterized.product(_OPTIMIZERS_UNDER_TEST)
+  def test_optimizers_accept_extra_args(self, opt_name, opt_kwargs):
+    opt = getattr(alias, opt_name)(**opt_kwargs)
+    # intentionally ommit: opt = base.with_extra_args_support(opt)
+    initial_params, _, objective = _setup_rosenbrock(jnp.float32)
+
+    @jax.jit
+    def step(params, state):
+      value, updates = jax.value_and_grad(objective)(params)
+      update_kwargs = {'unexpected_extra_args_your_optimizer_doesnt_expect': 1}
+      if opt_name in ['polyak_sgd']:
+        update_kwargs = {'value': value}
+      updates, state = opt.update(updates, state, params, **update_kwargs)
+      params = update.apply_updates(params, updates)
+      return params, state
+
+    params = initial_params
+    with self.subTest('Test that init works with extra values'):
+      state = opt.init(params)
+
+    with self.subTest('Test that update works with extra values'):
+      for _ in range(2):
+        params, state = step(params, state)
+
   @chex.all_variants
   @parameterized.product(_OPTIMIZERS_UNDER_TEST)
   def test_optimizers_can_be_wrapped_in_inject_hyperparams(

optax/_src/base.py

@@ -150,6 +150,8 @@ def __call__(
 class GradientTransformation(NamedTuple):
   # pylint: disable=line-too-long
   """A pair of pure functions implementing a gradient transformation.
+  
+  Prefer :class:`GradientTransformationExtraArgs` for new optimizers.
 
   Optax optimizers are all implemented as *gradient transformations*.
   A gradient transformation is defined to be a pair of pure functions, which

optax/_src/transform.py

@@ -975,8 +975,8 @@ def init_fn(params):
     del params
     return ScaleByScheduleState(count=jnp.zeros([], jnp.int32))
 
-  def update_fn(updates, state, params=None):
-    del params
+  def update_fn(updates, state, params=None, **extra_args):
+    del params, extra_args
     step_size = step_size_fn(state.count)
     updates = jax.tree.map(
         lambda g: jnp.array(step_size, dtype=g.dtype) * g, updates

optax/contrib/_common_test.py

@@ -213,6 +213,35 @@ def obj_fn(params):
 
 class ContribTest(chex.TestCase):
 
+  @parameterized.product(_ALL_OPTIMIZERS_UNDER_TEST, wrap=[True, False])
+  def test_optimizers_accept_extra_args(
+      self, opt_name, opt_kwargs, wrapper_name, wrapper_kwargs, wrap):
+    opt = _get_opt_factory(opt_name)(**opt_kwargs)
+    if wrap and wrapper_name is not None:
+      opt = _wrap_opt(opt, wrapper_name, wrapper_kwargs)
+    # intentionally ommit: opt = base.with_extra_args_support(opt)
+
+    initial_params, _, objective = _setup_rosenbrock(jnp.float32)
+
+    @jax.jit
+    def step(params, state):
+      value, updates = jax.value_and_grad(objective)(params)
+      update_kwargs = {'unexpected_extra_args_your_optimizer_doesnt_expect': 1}
+      if opt_name in ['momo', 'momo_adam', 'sgd']:
+        update_kwargs['value'] = value
+      if opt_name in ['sophia']:
+        update_kwargs['obj_fn'] = objective
+      updates, state = opt.update(updates, state, params, **update_kwargs)
+      params = update.apply_updates(params, updates)
+      return params, state
+
+    params = initial_params
+    state = opt.init(params)
+
+    with self.subTest('Test that update works with extra args'):
+      for _ in range(2):
+        params, state = step(params, state)
+
   @parameterized.product(
       _ALL_OPTIMIZERS_UNDER_TEST,
       target=(

optax/contrib/_dadapt_adamw.py

@@ -46,7 +46,7 @@ def dadapt_adamw(
     eps: float = 1e-8,
     estim_lr0: float = 1e-6,
     weight_decay: float = 0.0,
-) -> base.GradientTransformation:
+) -> base.GradientTransformationExtraArgs:
   """Learning rate free AdamW by D-Adaptation.
 
   Adapts the baseline learning rate of AdamW automatically by estimating the
@@ -91,7 +91,9 @@ def update_fn(
       updates: base.Updates,
       state: DAdaptAdamWState,
       params: Optional[base.Params] = None,
+      **extra_args,
   ) -> tuple[base.Updates, DAdaptAdamWState]:
+    del extra_args
     if params is None:
       raise ValueError(base.NO_PARAMS_MSG)
     count = state.count
@@ -141,4 +143,4 @@ def update_fn(
     )
     return p_update, new_state
 
-  return base.GradientTransformation(init_fn, update_fn)
+  return base.GradientTransformationExtraArgs(init_fn, update_fn)