flax.nnx

flax.nnx#

Experimental API. See the NNX page for more details.