LICENSE
README.md
pyproject.toml
setup.py
src/flax_extra/__init__.py
src/flax_extra/batch.py
src/flax_extra/console.py
src/flax_extra/random.py
src/flax_extra/util.py
src/flax_extra/version.py
src/flax_extra.egg-info/PKG-INFO
src/flax_extra.egg-info/SOURCES.txt
src/flax_extra.egg-info/dependency_links.txt
src/flax_extra.egg-info/not-zip-safe
src/flax_extra.egg-info/requires.txt
src/flax_extra.egg-info/top_level.txt
src/flax_extra/checkpoint/__init__.py
src/flax_extra/checkpoint/_best_checkpoint_file_reader.py
src/flax_extra/checkpoint/_best_checkpoint_file_writer.py
src/flax_extra/checkpoint/_checkpoint.py
src/flax_extra/checkpoint/_checkpoint_file.py
src/flax_extra/checkpoint/_checkpoint_file_reader.py
src/flax_extra/checkpoint/_checkpoint_file_writer.py
src/flax_extra/checkpoint/_summary.py
src/flax_extra/checkpoint/_summary_logger.py
src/flax_extra/checkpoint/_summary_writer.py
src/flax_extra/evaluation/__init__.py
src/flax_extra/evaluation/_loop.py
src/flax_extra/layer/__init__.py
src/flax_extra/layer/_lstm.py
src/flax_extra/model/__init__.py
src/flax_extra/model/_rnn.py
src/flax_extra/operator/__init__.py
src/flax_extra/operator/_shift_right.py
src/flax_extra/training/__init__.py
src/flax_extra/training/_loop.py