README.md
setup.py
trax/__init__.py
trax/history.py
trax/import_test.py
trax/jaxboard.py
trax/learning_rate.py
trax/learning_rate_test.py
trax/math.py
trax/rl_trainer.py
trax/shapes.py
trax/shapes_test.py
trax/trainer.py
trax/trainer_flags.py
trax/utils.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/layers/__init__.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/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/tracer.py
trax/layers/tracer_test.py
trax/layers/research/__init__.py
trax/layers/research/efficient_attention.py
trax/layers/research/efficient_attention_test.py
trax/layers/research/efficient_attention_v2.py
trax/layers/research/efficient_attention_v2_test.py
trax/layers/research/position_encodings.py
trax/math/__init__.py
trax/math/backend.py
trax/math/backend_test.py
trax/math/jax.py
trax/math/numpy.py
trax/math/tf.py
trax/models/__init__.py
trax/models/atari_cnn.py
trax/models/atari_cnn_test.py
trax/models/beam_search.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/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_test.py
trax/models/research/__init__.py
trax/models/research/position_lookup_transformer.py
trax/models/research/skipping_transformer.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/sgd.py
trax/optimizers/sm3.py
trax/rl/__init__.py
trax/rl/awr_trainer.py
trax/rl/awr_trainer_test.py
trax/rl/awr_utils.py
trax/rl/awr_utils_test.py
trax/rl/base_trainer.py
trax/rl/base_trainer_test.py
trax/rl/online_tune.py
trax/rl/online_tune_test.py
trax/rl/policy_based_trainer.py
trax/rl/ppo.py
trax/rl/ppo_test.py
trax/rl/ppo_trainer.py
trax/rl/ppo_trainer_test.py
trax/rl/serialization_utils.py
trax/rl/serialization_utils_test.py
trax/rl/simple.py
trax/rl/simple_test.py
trax/rl/simple_trainer.py
trax/rl/simple_trainer_test.py
trax/rl/simulated_env_problem.py
trax/rl/simulated_env_problem_test.py
trax/rl/space_serializer.py
trax/rl/space_serializer_test.py
trax/rl/trainers.py
trax/rl/envs/__init__.py
trax/rl/envs/async_trajectory_collector.py
trax/rl/envs/async_trajectory_collector_lib.py
trax/rl/envs/async_trajectory_collector_lib_test.py
trax/rl/envs/fake_env.py
trax/rl/envs/fake_env_test.py
trax/rl/envs/online_tune.py
trax/rl/envs/online_tune_env.py
trax/rl/envs/online_tune_env_test.py
trax/rl/envs/online_tune_rl_env.py
trax/rl/envs/online_tune_rl_env_test.py
trax/rl/envs/online_tune_test.py
trax/supervised/__init__.py
trax/supervised/inputs.py
trax/supervised/inputs_test.py
trax/supervised/mnist_test.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/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/array_creation.py
trax/tf_numpy/numpy/array_manipulation.py
trax/tf_numpy/numpy/array_methods.py
trax/tf_numpy/numpy/arrays.py
trax/tf_numpy/numpy/dtypes.py
trax/tf_numpy/numpy/logic.py
trax/tf_numpy/numpy/math.py
trax/tf_numpy/numpy/random.py
trax/tf_numpy/numpy/utils.py