LICENSE
README.md
setup.py
jax_models/__init__.py
jax_models.egg-info/PKG-INFO
jax_models.egg-info/SOURCES.txt
jax_models.egg-info/dependency_links.txt
jax_models.egg-info/requires.txt
jax_models.egg-info/top_level.txt
jax_models/activations/__init__.py
jax_models/activations/activations.py
jax_models/layers/__init__.py
jax_models/layers/attention.py
jax_models/layers/depthwise_separable_conv.py
jax_models/layers/drop.py
jax_models/layers/mask.py
jax_models/layers/mlp.py
jax_models/layers/patch_util.py
jax_models/layers/pool.py
jax_models/layers/squeeze_and_excite_layer.py
jax_models/models/__init__.py
jax_models/models/conv_mixer.py
jax_models/models/convnext.py
jax_models/models/masked_autoencoder.py
jax_models/models/mlp_mixer.py
jax_models/models/model_registry.py
jax_models/models/mpvit.py
jax_models/models/patchconvnet.py
jax_models/models/poolformer.py
jax_models/models/segformer.py
jax_models/models/swin_transformer.py