AUTHORS
LICENSE
README.md
setup.cfg
setup.py
flax/__init__.py
flax/config.py
flax/errors.py
flax/jax_utils.py
flax/py.typed
flax/serialization.py
flax/struct.py
flax/traceback_util.py
flax/traverse_util.py
flax/version.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/core/__init__.py
flax/core/axes_scan.py
flax/core/frozen_dict.py
flax/core/lift.py
flax/core/scope.py
flax/core/tracers.py
flax/core/variables.py
flax/core/nn/__init__.py
flax/core/nn/attention.py
flax/core/nn/linear.py
flax/core/nn/normalization.py
flax/core/nn/stochastic.py
flax/linen/__init__.py
flax/linen/activation.py
flax/linen/attention.py
flax/linen/combinators.py
flax/linen/dotgetter.py
flax/linen/dtypes.py
flax/linen/initializers.py
flax/linen/linear.py
flax/linen/module.py
flax/linen/normalization.py
flax/linen/partitioning.py
flax/linen/pooling.py
flax/linen/recurrent.py
flax/linen/stochastic.py
flax/linen/transforms.py
flax/metrics/__init__.py
flax/metrics/tensorboard.py
flax/optim/__init__.py
flax/optim/adabelief.py
flax/optim/adadelta.py
flax/optim/adafactor.py
flax/optim/adagrad.py
flax/optim/adam.py
flax/optim/base.py
flax/optim/dynamic_scale.py
flax/optim/lamb.py
flax/optim/lars.py
flax/optim/momentum.py
flax/optim/rmsprop.py
flax/optim/sgd.py
flax/optim/weight_norm.py
flax/testing/__init__.py
flax/testing/benchmark.py
flax/training/__init__.py
flax/training/checkpoints.py
flax/training/common_utils.py
flax/training/early_stopping.py
flax/training/lr_schedule.py
flax/training/prefetch_iterator.py
flax/training/train_state.py