Treat integer hyperparameters as static in inject_hyperparams (fixes #412)

optax/schedules/_inject.py

@@ -116,7 +116,10 @@ def inject_hyperparams(
     static_args: a string or iterable of strings specifying which callable
       parameters are not schedules. inject_hyperparams treats all callables as
       schedules by default, so if a hyperparameter is a non-schedule callable,
-      you must specify that using this argument.
+      you must specify that using this argument. Boolean and integer
+      hyperparameters are always treated as static (they are not injected),
+      since tracing them would break inner factories that use them for
+      structural control flow.
     hyperparam_dtype: Optional datatype override. If specified, all float
       hyperparameters will be cast to this type.
 
@@ -151,14 +154,19 @@ def wrapped_transform(
 
     sched_hps, numeric_hps, other_hps = {}, {}, {}
     for name, value in bound_arguments.arguments.items():
-      if name in static_args or isinstance(value, bool):
+      # Python `bool`/`int` values are treated as static. Turning them into
+      # traced arrays would break inner factories that use them for structural
+      # control flow (e.g. `min_dim_size_to_factor` in `adafactor` or
+      # `memory_size` in `lbfgs`), causing a `TracerBoolConversionError` when
+      # the resulting transform is jitted. Note `bool` is a subclass of `int`.
+      if name in static_args or isinstance(value, int):
         other_hps[name] = value
       elif isinstance(value, base.StatefulSchedule):
         sched_hps[name] = value
       elif callable(value):
         # pyrefly: ignore[bad-argument-type]
         sched_hps[name] = WrappedSchedule(value)
-      elif isinstance(value, (int, float, jax.Array, np.ndarray)):
+      elif isinstance(value, (float, jax.Array, np.ndarray)):
         numeric_hps[name] = value
       else:
         other_hps[name] = value
@@ -238,7 +246,10 @@ def inject_stateful_hyperparams(
     static_args: a string or iterable of strings specifying which callable
       parameters are not schedules. inject_hyperparams treats all callables as
       schedules by default, so if a hyperparameter is a non-schedule callable,
-      you must specify that using this argument.
+      you must specify that using this argument. Boolean and integer
+      hyperparameters are always treated as static (they are not injected),
+      since tracing them would break inner factories that use them for
+      structural control flow.
     hyperparam_dtype: Optional datatype override. If specified, all float
       hyperparameters will be cast to this type.
 

optax/schedules/_inject_test.py

@@ -27,6 +27,7 @@
 from optax import schedules
 from optax import transforms
 
+from optax._src import alias
 from optax._src import base
 from optax._src import test_utils
 from optax._src import transform
@@ -159,6 +160,29 @@ def test_numeric_static_args(self, static_args):
 
     assert not set(state.hyperparams.keys()).intersection(set(static_args))
 
+  def test_integer_hyperparams_are_static(self):
+    """Integer hyperparameters are not injected (kept static)."""
+    optim = schedules.inject_hyperparams(alias.adafactor)(learning_rate=0.1)
+
+    params = jnp.ones((4, 4))
+    state = jax.jit(optim.init)(params)
+
+    # `min_dim_size_to_factor` is a structural int and must not be injected.
+    self.assertNotIn('min_dim_size_to_factor', state.hyperparams)
+    self.assertIn('learning_rate', state.hyperparams)
+
+  @parameterized.named_parameters(
+      ('adafactor', alias.adafactor),  # int `min_dim_size_to_factor`
+      ('lbfgs', alias.lbfgs),  # int `memory_size`
+  )
+  def test_jit_init_with_integer_hyperparams(self, opt_factory):
+    """Optimizers with integer hyperparams jit without `static_args` (#412)."""
+    optim = schedules.inject_hyperparams(opt_factory)(learning_rate=0.1)
+
+    params = jnp.ones((4, 4))
+    # Prior to treating ints as static this raised a TracerBoolConversionError.
+    jax.jit(optim.init)(params)
+
   def test_prng_key_not_hyperparameter(self):
     """Check that random.key can be handled by :func:``inject_hyperparams``."""