Added (*) in "axlearn/common/flash_attention/layer_test.py" L529 to

unpack Tuple before converting it to PartitionSpec. This fixes the
issue of having:

AssertionError: P('data', None) != P(('data', None),)

Author: QuinnMMcGarry

axlearn/common/flash_attention/layer_test.py

@@ -526,7 +526,7 @@ def as_partition_spec(pytree: CompositeAttentionBias) -> PartitionSpec:
             spec = test_layer._logit_biases_spec(segment_ids)  # pylint: disable=protected-access
             spec = as_partition_spec(spec)
             self.assertIsInstance(spec, PartitionSpec)
-            self.assertEqual(spec, as_partition_spec(test_layer.config.mha_dim_to_partition_spec["btnh"][:2]))
+            self.assertEqual(spec, jax.P(*test_layer.config.mha_dim_to_partition_spec["btnh"][:2]))
 
     @parameterized.product(
         _TEST_CONFIGS,

pyproject.toml

@@ -25,8 +25,8 @@ core = [
     "chex==0.1.91",
     "etils[epy]==1.13.0",  # pin to ensure importlib-resources is installed (needed by tensorflow-datasets)
     "importlab==0.8.1",  # breaks pytype on 0.8
-    "jax>=0.9.1.dev20260130,<0.9.2",
-    "jaxlib>=0.9.1.dev20260130,<0.9.2",
+    "jax==0.10.0",
+    "jaxlib>=0.10.0",
     "ml-dtypes==0.5.3",
     "msgpack==1.1.0",  # for checkpointing.
     "nltk==3.7",  # for text preprocessing
@@ -108,7 +108,7 @@ gcp = [
 # Note: Specify -f https://storage.googleapis.com/jax-releases/libtpu_releases.html during install.
 tpu = [
     "axlearn[gcp]",
-    "jax[tpu]>=0.9.1.dev20260130,<0.9.2",  # must be >=0.4.19 for compat with v5p.
+    "jax[tpu]==0.10.0",  # must be >=0.4.19 for compat with v5p.
     "pathwaysutils==0.1.2",  # For JAX+Pathways single-controller accelerator coordinator.
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
@@ -134,7 +134,7 @@ dataflow = [
 # GPU custom kernel dependency.
 gpu = [
     "triton>=2.1.0",
-    "jax[cuda12]>=0.9.1.dev20260130,<0.9.2",
+    "jax[cuda12]==0.10.0",
     "nvidia-ml-py==12.560.30",
     # pin nccl version, otherwise jax[cuda12] will pull latest version
     "nvidia-nccl-cu12==2.27.5",