Merge pull request #31439 from chapman20j:abstract_device

PiperOrigin-RevId: 801031638

Author: Google-ML-Automation

jax/sharding.py

@@ -27,6 +27,7 @@
 )
 from jax._src.mesh import (
     Mesh as Mesh,
+    AbstractDevice as AbstractDevice,
     AbstractMesh as AbstractMesh,
     AxisType as AxisType,
     get_abstract_mesh as get_abstract_mesh,

tests/pjit_test.py

@@ -41,7 +41,7 @@
 from jax.lax import with_sharding_constraint
 from jax._src import prng
 from jax.sharding import (PartitionSpec as P, Mesh, auto_axes, explicit_axes,
-                          reshard)
+                          reshard, AbstractDevice)
 from jax.experimental import multihost_utils
 from jax._src.shard_map import shard_map
 from jax._src.compilation_cache import is_persistent_cache_enabled
@@ -53,7 +53,7 @@
 from jax._src.sharding_impls import (
     AUTO, UNSPECIFIED, NamedSharding, GSPMDSharding,
     SingleDeviceSharding, parse_flatten_op_sharding)
-from jax._src.mesh import use_abstract_mesh, AbstractDevice
+from jax._src.mesh import use_abstract_mesh
 from jax._src.pjit import pjit, _pjit_lower
 from jax._src.layout import Format, Layout as DLL
 from jax._src.named_sharding import DuplicateSpecError