flax

[flax:examples:lm1b_nnx] Update example to work internally. #jax-fixit.

#5125

Merged