Access thread_resources via jax.interpreters.pxla instead of jax.expe…

…rimental.maps

The maps submodule is deprecated and will be removed soon.

PiperOrigin-RevId: 617511370

flax/core/meta.py

@@ -28,7 +28,7 @@
 from flax import errors, struct
 from flax.typing import LogicalNames
 import jax
-from jax.experimental import maps
+from jax.interpreters import pxla
 
 A = TypeVar('A')
 B = TypeVar('B')
@@ -178,9 +178,9 @@ def inner_update(c, v):
 
 
 def _global_mesh_defined() -> bool:
-  """Checks if global xmap/pjit mesh resource environment is defined."""
-  maps_env = maps.thread_resources.env
-  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
+  """Checks if global mesh resource environment is defined."""
+  env = pxla.thread_resources.env
+  return env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
 
 
 class Partitioned(struct.PyTreeNode, AxisMetadata[A]):

flax/experimental/nnx/nnx/spmd.py

@@ -16,7 +16,7 @@
 import typing as tp
 
 import jax
-from jax.experimental import maps
+from jax.interpreters import pxla
 from jax.sharding import Mesh, PartitionSpec
 
 from flax.experimental.nnx.nnx import variables
@@ -126,9 +126,9 @@ def get_named_sharding(tree: A, mesh: jax.sharding.Mesh) -> A:
 
 
 def _global_mesh_defined() -> bool:
-  """Checks if global xmap/pjit mesh resource environment is defined."""
-  maps_env = maps.thread_resources.env
-  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
+  """Checks if global mesh resource environment is defined."""
+  env = pxla.thread_resources.env
+  return env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
 
 
 def _with_sharding_constraint(

flax/linen/spmd.py

@@ -34,7 +34,7 @@
 
 import jax
 from jax import lax
-from jax.experimental import maps
+from jax.interpreters import pxla
 
 from flax import struct
 from flax.core import meta
@@ -209,9 +209,9 @@ def logical_to_mesh_sharding(
 
 
 def _global_mesh_defined() -> bool:
-  """Checks if global xmap/jit mesh resource environment is defined."""
-  maps_env = maps.thread_resources.env
-  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
+  """Checks if global mesh resource environment is defined."""
+  env = pxla.thread_resources.env
+  return env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
 
 
 class RulesFallback(enum.Enum):