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

[ArneshBanerjee]

Summary

optax.inject_hyperparams converts every numeric argument into a traced array so it can be scheduled/overridden at runtime. Boolean arguments are already special-cased as static, because a traced boolean can't be used in Python control flow.

Integer arguments have the identical problem, but currently fall through to the "numeric" branch and get traced. Several optimizers use integer arguments for structural decisions:

• min_dim_size_to_factor in adafactor (optax/_src/factorized.py:55 → if shape[...] < min_dim_size_to_factor)
• memory_size in lbfgs (optax/_src/transform.py:1717 → if memory_size < 1)

As a result, jitting these fails with TracerBoolConversionError unless the user manually passes static_args=(...):

importjax, jax.numpyasjnp, optaxopt=optax.inject_hyperparams(optax.adafactor)(learning_rate=0.1)
jax.jit(opt.init)(jnp.ones((4, 4))) # TracerBoolConversionError

This is exactly what #412 asks to fix ("all optimizers wrapped in inject_hyperparams can be jit compiled without any additional static_args").

Fix

Since bool is a subclass of int, extending the existing static-value check from bool to int generalizes the current behavior and resolves both optimizers with no new special cases. Integer hyperparameters cannot be meaningfully scheduled under jit anyway (a schedule returns floats, and structural ints break control flow when traced).

Verification

• inject_hyperparams(optax.adafactor) and inject_hyperparams(optax.lbfgs) now jax.jit out of the box; the previously-required static_args workaround is no longer needed.
• Added regression tests in optax/schedules/_inject_test.py.
• pytest optax/schedules/ optax/_src/alias_test.py → 634 passed, 70 skipped, no regressions.

Fixes#412.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.