Expose get_concrete_mesh to users through jax.sharding.

Yun Ni · Google-ML-Automation · commit df8ce2014d7b · 2025-10-01T11:05:36.000-07:00
This allows get_concrete_mesh() to retrieve the mesh set by jax.set_mesh(), which makes the API more cohesive.

PiperOrigin-RevId: 813819319
diff --git a/jax/sharding.py b/jax/sharding.py
@@ -31,6 +31,7 @@
     AbstractMesh as AbstractMesh,
     AxisType as AxisType,
     get_abstract_mesh as get_abstract_mesh,
+    get_concrete_mesh as get_concrete_mesh,
     use_abstract_mesh as use_abstract_mesh,
 )
 

Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,7 @@`
`31`	`31`	`AbstractMesh as AbstractMesh,`
`32`	`32`	`AxisType as AxisType,`
`33`	`33`	`get_abstract_mesh as get_abstract_mesh,`
	`34`	`+ get_concrete_mesh as get_concrete_mesh,`
`34`	`35`	`use_abstract_mesh as use_abstract_mesh,`
`35`	`36`	`)`
`36`	`37`