Trax is meant as a successor to Tensor2Tensor. From my understanding of what Tensor2Tensor was (and thus what Trax is): Basically, at it's core, it's a librar

Trax is meant as a successor to Tensor2Tensor. From my understanding of what Tensor2Tensor was (and thus what Trax is): Basically, at it's core, it's a library of models, datasets, and utilities that are meant to be clear and work together nicely. This is meant to be good for education, as you have an easy and unified way to run SOTA trained models on a wide variety of datasets. This is also meant to be good for research, as much of research involves building on top of other people's models, and having a library of models with shared API (as opposed to the fragmented wilderness that is public research code) is meant to facilitate that. I believe that it also ended up being used in industry as well, simply as a reliable source of pre-trained models. PS: As Jax did not have an "official" neural network library until very recently, this also served as a neural network library for Jax.