Trax helps you understand deep learning. We start with basic maths and go through layers, models, supervised and reinforcement learning. We get to advanced deep learning results, including recent papers and state-of-the-art models.

Trax is a successor to the Tensor2Tensor library and is actively used and maintained by researchers and engineers within the Google Brain team and a community of users. We're eager to collaborate with you too, so feel free to open an issue on GitHub or send along a pull request (see our contribution doc).

See our example layers in a TPU/GPU-backed colab notebook at Trax Demo

python -m trax.trainer \
  --dataset=mnist \
  --model=MLP \
  --config="train.steps=1000"

python -m trax.trainer \
  --config_file=$PWD/trax/configs/resnet50_imagenet_8gb.gin

python -m trax.trainer \
  --config_file=transformer_lm1b_8gb.gin