Add Transformer Engine support to Paxml

[ashors1]

Adds support for NVIDIA's Transformer Engine. TE can be enabled by setting the environment variable ENABLE_TE=1. For more details about running Pax with Transformer Engine, refer to the JAX Toolbox README.