Fix: Fixed UT jax.make_mesh should take jax.sharding.AxisType.Auto in tests

Author: camiloCienet

axlearn/common/kv_cache/kv_cache_test.py

@@ -5,6 +5,7 @@
 import jax
 import jax.numpy as jnp
 import pytest
+from jax.sharding import AxisType
 from absl.testing import absltest, parameterized
 from jax.sharding import PartitionSpec
 
@@ -173,7 +174,7 @@ def test_init_states_kv_partition_spec(self):
         )
         shape = KVCache.Shape(batch_size=batch, kv_len=kv_len, num_kv_heads=heads, per_head_dim=dim)
 
-        with jax.make_mesh((4, 2), ("data", "model")):
+        with jax.make_mesh((4, 2), ("data", "model"), axis_types=(AxisType.Auto, AxisType.Auto)):
 
             @jax.jit
             def f():

axlearn/common/kv_cache/sliding_window_kv_cache_test.py

@@ -4,6 +4,7 @@
 
 import jax
 import jax.numpy as jnp
+from jax.sharding import AxisType
 import pytest
 from absl.testing import absltest, parameterized
 from jax.sharding import PartitionSpec
@@ -131,7 +132,7 @@ def test_init_states_kv_partition_spec(self):
             batch_size=batch, kv_len=32, num_kv_heads=heads, per_head_dim=dim
         )
 
-        with jax.make_mesh((4, 2), ("data", "model")):
+        with jax.make_mesh((4, 2), ("data", "model"), axis_types=(AxisType.Auto, AxisType.Auto)):
 
             @jax.jit
             def f():