Fix API usage of jax.make_mesh (#1051)

Signed-off-by: Lihao Ran <[email protected]>

Author: Lumosis

examples/disagg/multi_proc_per_worker.py

@@ -80,7 +80,11 @@ def get_uuid(pid: int) -> int:
 
 def get_mesh() -> Mesh:
     sharding_size = jax.device_count()
-    return jax.make_mesh((sharding_size, ), ("model", ))
+    return jax.make_mesh(
+        (sharding_size, ),
+        ("model", ),
+        axis_types=(jax.sharding.AxisType.Auto, ) * len(("model", )),
+    )
 
 
 def get_kv_shape(mesh: Mesh, per_shard: bool = False) -> tuple[int, ...]:

examples/disagg/single_proc_per_worker.py

@@ -74,7 +74,11 @@ def get_uuid() -> int:
 
 def get_mesh() -> Mesh:
     sharding_size = jax.device_count()
-    return jax.make_mesh((sharding_size, ), ("model", ))
+    return jax.make_mesh(
+        (sharding_size, ),
+        ("model", ),
+        axis_types=(jax.sharding.AxisType.Auto, ) * len(("model", )),
+    )
 
 
 def get_kv_spec(mesh: Mesh) -> list[int]: