flax

Update optimizer.py to support masked variable from optax.

#4904

Merged