# Pure Optax (without NNX)
import optax

optimizer = optax.adam(learning_rate=0.001)
opt_state = optimizer.init(params)

# In training loop:
updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)

        __
        __