diff --git a/optax/schedules/_inject.py b/optax/schedules/_inject.py index 6688587cf..20bd0bde2 100644 --- a/optax/schedules/_inject.py +++ b/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. diff --git a/optax/schedules/_inject_test.py b/optax/schedules/_inject_test.py index 301636abd..e32168bf4 100644 --- a/optax/schedules/_inject_test.py +++ b/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``."""