README.md
setup.py
flax/__init__.py
flax/jax_utils.py
flax/optim.py
flax/serialization.py
flax/struct.py
flax/traverse_util.py
flax.egg-info/PKG-INFO
flax.egg-info/SOURCES.txt
flax.egg-info/dependency_links.txt
flax.egg-info/not-zip-safe
flax.egg-info/requires.txt
flax.egg-info/top_level.txt
flax/metrics/__init__.py
flax/metrics/tensorboard.py
flax/nn/__init__.py
flax/nn/activation.py
flax/nn/attention.py
flax/nn/base.py
flax/nn/initializers.py
flax/nn/linear.py
flax/nn/normalization.py
flax/nn/pooling.py
flax/nn/recurrent.py
flax/nn/stochastic.py
flax/nn/utils.py
flax/training/__init__.py
flax/training/checkpoints.py
flax/training/common_utils.py
flax/training/lr_schedule.py
flax/training/prefetch_iterator.py