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/py.typed
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/combinator/__init__.py
src/flax_extra/combinator/_concatenate.py
src/flax_extra/data/__init__.py
src/flax_extra/data/_bytes_tokenizer.py
src/flax_extra/evaluation/__init__.py
src/flax_extra/evaluation/_loop.py
src/flax_extra/layer/__init__.py
src/flax_extra/layer/_attention.py
src/flax_extra/layer/_decoding.py
src/flax_extra/layer/_embedding_decoding.py
src/flax_extra/layer/_encoding.py
src/flax_extra/layer/_feedforward.py
src/flax_extra/layer/_fourier_positional_encoding.py
src/flax_extra/layer/_lstm.py
src/flax_extra/layer/_multimodal_decoding.py
src/flax_extra/layer/_multimodal_encoding.py
src/flax_extra/layer/_multimodal_positional_encoding.py
src/flax_extra/layer/_trainable_positional_encoding.py
src/flax_extra/layer/_trainable_positional_masking.py
src/flax_extra/layer/_trainable_positional_padding.py
src/flax_extra/layer/io.py
src/flax_extra/model/__init__.py
src/flax_extra/model/perceiver/__init__.py
src/flax_extra/model/perceiver/_cross_attention_block.py
src/flax_extra/model/perceiver/_decoder.py
src/flax_extra/model/perceiver/_encoder.py
src/flax_extra/model/perceiver/_io.py
src/flax_extra/model/perceiver/_processor.py
src/flax_extra/model/perceiver/_self_attention_block.py
src/flax_extra/model/rnn/__init__.py
src/flax_extra/model/rnn/_lm.py
src/flax_extra/operator/__init__.py
src/flax_extra/operator/_rearrange.py
src/flax_extra/operator/_reshape.py
src/flax_extra/operator/_shift_right.py
src/flax_extra/training/__init__.py
src/flax_extra/training/_loop.py