flax-extra (pypi): The package provides extra flexibility to Flax using ideas originated at Trax
Keywords: flax|jax|trax