Don't require installing trax dependencies (jax, jaxlib) since they

aren't available for Windows.

PiperOrigin-RevId: 242481869

Author: afrozenator

oss_scripts/oss_pip_install.sh

@@ -22,6 +22,6 @@ t2t-trainer --registry_help 2>&1 >/dev/null
 t2t-datagen 2>&1 | grep translate_ende 2>&1 >/dev/null && echo passed
 
 # Then install the test dependencies
-pip install -q -e .[tests,allen]
+pip install -q -e .[tests,allen,trax]
 # Make sure to install the atari extras for gym
 pip install "gym[atari]"

setup.py

@@ -43,8 +43,6 @@
         'gunicorn',
         'gym',
         'h5py',
-        'jax',
-        'jaxlib',
         'kfac',
         'mesh-tensorflow',
         'numpy',
@@ -77,6 +75,10 @@
             # explicit pip install gym[atari] for the tests.
             # 'gym[atari]',
         ],
+        'trax': [
+            'jax',
+            'jaxlib',
+        ],
         'allen': ['Pillow==5.1.0', 'pandas==0.23.0'],
     },
     classifiers=[