Disable e2e tests when jax_pmap_shmap_merge=True.

`trainer._multi_device_update_fn` uses `jax.pmap` and when `jax_pmap_shmap_merge=True`, `jax.pmap` requires inputs be explicitly sharded as the underlying `jax.jit` expects.

This would need to be fixed for when `jax_pmap_shmap_merge=True` by default.

PiperOrigin-RevId: 811399636

Author: danielsuo

trax/models/reformer/reformer_e2e_test.py

@@ -19,6 +19,7 @@
 
 from absl.testing import absltest
 import gin
+import jax
 
 from trax import test_utils
 from trax.models.reformer import reformer  # pylint: disable=unused-import
@@ -36,6 +37,11 @@ def setUp(self):
     super().setUp()
     gin.clear_config()
     gin.add_config_file_search_path(_CONFIG_DIR)
+    # NOTE(dsuo): trainer._multi_device_update_fn uses `jax.pmap` and when
+    # `jax_pmap_shmap_merge=True`, `jax.pmap` requires inputs be explicitly
+    # sharded as the underlying `jax.jit` expects. This would need to be fixed.
+    if jax.config.jax_pmap_shmap_merge:
+      self.skipTest('Skipping test because jax_pmap_shmap_merge is enabled.')
     test_utils.ensure_flag('test_tmpdir')
 
   def test_reformer_wmt_ende(self):

trax/models/research/terraformer_e2e_test.py

@@ -19,6 +19,7 @@
 
 from absl.testing import absltest
 import gin
+import jax
 
 from trax import test_utils
 from trax.models.research import terraformer  # pylint: disable=unused-import
@@ -36,6 +37,11 @@ def setUp(self):
     super().setUp()
     gin.clear_config()
     gin.add_config_file_search_path(_CONFIG_DIR)
+    # NOTE(dsuo): trainer._multi_device_update_fn uses `jax.pmap` and when
+    # `jax_pmap_shmap_merge=True`, `jax.pmap` requires inputs be explicitly
+    # sharded as the underlying `jax.jit` expects. This would need to be fixed.
+    if jax.config.jax_pmap_shmap_merge:
+      self.skipTest('Skipping test because jax_pmap_shmap_merge is enabled.')
     test_utils.ensure_flag('test_tmpdir')
 
   def test_terraformer_wmt_ende(self):

trax/supervised/trainer_lib_test.py

@@ -156,6 +156,11 @@ def setUp(self):
     super().setUp()
     test_utils.ensure_flag('test_tmpdir')
     self._old_is_allow_float64 = tf_np.is_allow_float64()
+    # NOTE(dsuo): trainer._multi_device_update_fn uses `jax.pmap` and when
+    # `jax_pmap_shmap_merge=True`, `jax.pmap` requires inputs be explicitly
+    # sharded as the underlying `jax.jit` expects. This would need to be fixed.
+    if jax.config.jax_pmap_shmap_merge:
+      self.skipTest('Skipping test because jax_pmap_shmap_merge is enabled.')
     tf_np.set_allow_float64(False)
 
   def tearDown(self):