# With NNX
optimizer = nnx.Optimizer(model, optax.adam(0.001), wrt=nnx.Param)
optimizer.update(model, grads)  # Handles everything internally

        __
        __