from flax import nnx
import optax

# Create the Optax optimizer
tx = optax.adam(learning_rate=0.001)

# Wrap it with nnx.Optimizer
optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)

        __
        __