LICENSE
README.md
setup.py
docs/conf.py
examples/denoising_diffusion/data_loader.py
examples/denoising_diffusion/model.py
examples/denoising_diffusion/train.py
examples/notebooks/pretrained_resnet18.py
examples/transformer/data.py
examples/transformer/model.py
examples/transformer/train.py
examples/wave_gru/data_loader.py
examples/wave_gru/model.py
examples/wave_gru/train.py
pax/__init__.py
pax/nets.py
pax/nn.py
pax/utils.py
pax/_src/__init__.py
pax/_src/utils.py
pax/_src/core/__init__.py
pax/_src/core/base.py
pax/_src/core/graph_module.py
pax/_src/core/mixed_precision.py
pax/_src/core/module.py
pax/_src/core/module_and_value.py
pax/_src/core/pure.py
pax/_src/core/rng.py
pax/_src/core/threading_local.py
pax/_src/core/transforms.py
pax/_src/core/utility_modules.py
pax/_src/core/utils.py
pax/_src/nets/__init__.py
pax/_src/nets/resnet.py
pax/_src/nets/transformer.py
pax/_src/nn/__init__.py
pax/_src/nn/attention.py
pax/_src/nn/batch_norm.py
pax/_src/nn/conv.py
pax/_src/nn/dropout.py
pax/_src/nn/ema.py
pax/_src/nn/embed.py
pax/_src/nn/group_norm.py
pax/_src/nn/identity.py
pax/_src/nn/lambda_module.py
pax/_src/nn/layer_norm.py
pax/_src/nn/linear.py
pax/_src/nn/pool.py
pax/_src/nn/recurrent.py
pax/_src/nn/rng_seq.py
pax/_src/nn/sequential.py
pax/experimental/__init__.py
pax/experimental/graph.py
pax3.egg-info/PKG-INFO
pax3.egg-info/SOURCES.txt
pax3.egg-info/dependency_links.txt
pax3.egg-info/not-zip-safe
pax3.egg-info/requires.txt
pax3.egg-info/top_level.txt