README.md
setup.py
trax/__init__.py
trax/import_test.py
trax/jaxboard.py
trax/predict_drop.py
trax/rl_trainer.py
trax/shapes.py
trax/shapes_test.py
trax/test_utils.py
trax/trainer.py
trax/trainer_flags.py
trax/trax2keras.py
trax/trax2keras_test.py
trax.egg-info/PKG-INFO
trax.egg-info/SOURCES.txt
trax.egg-info/dependency_links.txt
trax.egg-info/requires.txt
trax.egg-info/top_level.txt
trax/data/__init__.py
trax/data/debug_data_pipeline.py
trax/data/inputs.py
trax/data/inputs_test.py
trax/data/text_encoder.py
trax/data/text_encoder_build_subword.py
trax/data/text_encoder_test.py
trax/data/tf_inputs.py
trax/data/tf_inputs_test.py
trax/data/tokenizer.py
trax/data/tokenizer_test.py
trax/fastmath/__init__.py
trax/fastmath/jax.py
trax/fastmath/numpy.py
trax/fastmath/ops.py
trax/fastmath/ops_test.py
trax/fastmath/tf.py
trax/layers/__init__.py
trax/layers/acceleration.py
trax/layers/acceleration_test.py
trax/layers/activation_fns.py
trax/layers/activation_fns_test.py
trax/layers/assert_shape.py
trax/layers/assert_shape_test.py
trax/layers/attention.py
trax/layers/attention_test.py
trax/layers/base.py
trax/layers/base_test.py
trax/layers/combinators.py
trax/layers/combinators_test.py
trax/layers/convolution.py
trax/layers/convolution_test.py
trax/layers/core.py
trax/layers/core_test.py
trax/layers/deconvolution.py
trax/layers/deconvolution_test.py
trax/layers/initializers.py
trax/layers/initializers_test.py
trax/layers/metrics.py
trax/layers/metrics_test.py
trax/layers/normalization.py
trax/layers/normalization_test.py
trax/layers/pooling.py
trax/layers/pooling_test.py
trax/layers/reversible.py
trax/layers/reversible_test.py
trax/layers/rnn.py
trax/layers/rnn_test.py
trax/layers/test_utils.py
trax/layers/test_utils_test.py
trax/layers/research/__init__.py
trax/layers/research/efficient_attention.py
trax/layers/research/efficient_attention_test.py
trax/layers/research/position_encodings.py
trax/layers/research/position_encodings_test.py
trax/layers/research/rel_attention.py
trax/layers/research/rel_attention_test.py
trax/layers/research/resampling.py
trax/layers/research/rotary_positional_embedding.py
trax/layers/research/rotary_positional_embedding_test.py
trax/layers/research/sparsity.py
trax/layers/research/sparsity_test.py
trax/models/__init__.py
trax/models/atari_cnn.py
trax/models/atari_cnn_test.py
trax/models/mlp.py
trax/models/mlp_test.py
trax/models/neural_gpu.py
trax/models/neural_gpu_test.py
trax/models/resnet.py
trax/models/resnet_test.py
trax/models/rl.py
trax/models/rl_test.py
trax/models/rnn.py
trax/models/rnn_test.py
trax/models/transformer.py
trax/models/transformer_test.py
trax/models/reformer/__init__.py
trax/models/reformer/reformer.py
trax/models/reformer/reformer_e2e_test.py
trax/models/reformer/reformer_memory_test.py
trax/models/reformer/reformer_test.py
trax/models/research/__init__.py
trax/models/research/bert.py
trax/models/research/configurable_transformer.py
trax/models/research/configurable_transformer_test.py
trax/models/research/hourglass.py
trax/models/research/hourglass_test.py
trax/models/research/layerdrop_transformer.py
trax/models/research/layerdrop_transformer_test.py
trax/models/research/predict_terraformer.py
trax/models/research/rezero.py
trax/models/research/rezero_test.py
trax/models/research/rse.py
trax/models/research/rse_test.py
trax/models/research/terraformer.py
trax/models/research/terraformer_e2e_test.py
trax/models/research/terraformer_memory_test.py
trax/models/research/terraformer_oom_test.py
trax/models/research/terraformer_test.py
trax/models/research/transformer2.py
trax/models/research/transformer2_test.py
trax/optimizers/__init__.py
trax/optimizers/adafactor.py
trax/optimizers/adam.py
trax/optimizers/base.py
trax/optimizers/momentum.py
trax/optimizers/optimizers_test.py
trax/optimizers/rms_prop.py
trax/optimizers/sm3.py
trax/optimizers/trainer.py
trax/optimizers/trainer_test.py
trax/rl/__init__.py
trax/rl/actor_critic.py
trax/rl/actor_critic_joint.py
trax/rl/actor_critic_joint_test.py
trax/rl/actor_critic_test.py
trax/rl/advantages.py
trax/rl/advantages_test.py
trax/rl/atari_test.py
trax/rl/distributions.py
trax/rl/distributions_test.py
trax/rl/normalization.py
trax/rl/normalization_test.py
trax/rl/policy_tasks.py
trax/rl/rl_layers.py
trax/rl/serialization_utils.py
trax/rl/serialization_utils_test.py
trax/rl/space_serializer.py
trax/rl/space_serializer_test.py
trax/rl/task.py
trax/rl/task_test.py
trax/rl/training.py
trax/rl/training_test.py
trax/rl/value_tasks.py
trax/rl/value_tasks_test.py
trax/rl/envs/__init__.py
trax/rl/envs/data_envs.py
trax/rl/envs/data_envs_test.py
trax/supervised/__init__.py
trax/supervised/callbacks.py
trax/supervised/callbacks_test.py
trax/supervised/decoding.py
trax/supervised/decoding_test.py
trax/supervised/decoding_timing_test.py
trax/supervised/history.py
trax/supervised/history_test.py
trax/supervised/lr_schedules.py
trax/supervised/lr_schedules_test.py
trax/supervised/mnist_test.py
trax/supervised/pretrain_finetune.py
trax/supervised/trainer_lib.py
trax/supervised/trainer_lib_test.py
trax/supervised/training.py
trax/supervised/training_test.py
trax/tf_numpy/__init__.py
trax/tf_numpy/public_symbol_test.py
trax/tf_numpy/extensions/__init__.py
trax/tf_numpy/extensions/extensions.py
trax/tf_numpy/extensions/extensions_test.py
trax/tf_numpy/numpy/__init__.py
trax/tf_numpy/numpy_impl/__init__.py
trax/tf_numpy/numpy_impl/array_ops.py
trax/tf_numpy/numpy_impl/arrays.py
trax/tf_numpy/numpy_impl/dtypes.py
trax/tf_numpy/numpy_impl/math_ops.py
trax/tf_numpy/numpy_impl/random.py
trax/tf_numpy/numpy_impl/utils.py