Add init Trax research models.
Add init Trax research models.

PiperOrigin-RevId: 247220362

tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor