inject_hyperparams for Runtime LR

Optax's inject_hyperparams wraps an optimizer constructor so numeric hyperparameters can be stored in its state or supplied as schedules. This exercise uses the wrapper to take one fresh Adam step at a supplied learning rate.

Wrap the constructor

Create the optimizer with:

optax.inject_hyperparams(optax.adam)(learning_rate=lr)

The first call wraps the constructor. The second supplies its arguments. Use learning_rate, which is Adam's argument name; lr is only the name of this exercise's input.

Initialize the optimizer from params, calculate an update using grads, and apply it. Leave Adam's other settings at their defaults. The first-step rule, in exact arithmetic, is:

new_params = params - lr * grads / (abs(grads) + 1e-8)

For gradients much larger in magnitude than 1e-8, each parameter moves by about lr against the gradient's sign. Tiny gradients receive smaller updates, and zero gradients produce no change.

What injection does and does not change

The wrapper exposes numeric hyperparameters in optimizer state and can evaluate scheduled values as updates proceed. This single-step task uses a constant supplied value and restarts Adam's state on each call. It does not carry a training history or implement a learning-rate schedule.

Injection is not a universal requirement for changing values under JIT. A scalar passed as a dynamic JAX argument can change without retracing when its shape and dtype stay compatible. Static arguments and values captured by a function behave differently. Whether a new compilation is needed depends on how the value enters the compiled function.