make jax_remove_size_one_mesh_axis_from_type param setting in try block, todo: remove this after updating jax. Configure ICI data parallelism for decoupled mode

Author: gulsumgudukbay

tests/attention_test_util.py

@@ -25,6 +25,8 @@
 from MaxText import max_utils
 from MaxText import maxtext_utils
 from MaxText import pyconfig
+from MaxText.gcloud_stub import is_decoupled
+from maxtext.tests.test_utils import get_test_config_path
 from MaxText.common_types import AttentionType, DECODING_ACTIVE_SEQUENCE_INDICATOR, EP_AS_CONTEXT, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, ShardMode
 from MaxText.globals import MAXTEXT_PKG_DIR
 from MaxText.layers.attention_mla import MLA
@@ -52,10 +54,23 @@ class MLATestBase(parameterized.TestCase):
   def setUp(self):
     """Initializes the configuration for each test"""
     super().setUp()
-    jax.config.update("jax_remove_size_one_mesh_axis_from_type", True)
+    config_args = dict(self.config_arguments)
+    if is_decoupled():  # TODO(gulsumgudukbay): remove this after jax is updated.
+      # Older/newer JAX versions may not recognize this flag; ignore if absent.
+      try:
+        jax.config.update("jax_remove_size_one_mesh_axis_from_type", True)
+      except AttributeError:
+        pass
+      # In decoupled mode, adapt mesh/ICI parallelism to local devices so
+      # fill_unspecified_mesh_axes matches the available device count.
+      config_args.setdefault("mesh_axes", ["data"])
+      config_args.setdefault("ici_data_parallelism", -1)
+    else:
+      jax.config.update("jax_remove_size_one_mesh_axis_from_type", True)
+
     config = pyconfig.initialize(
-        [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")],
-        **self.config_arguments,
+        [sys.argv[0], get_test_config_path()],
+        **config_args,
     )
     self.cfg = config
     self.rng = jax.random.PRNGKey(0)