Repository: jax-ml/jax
Pull request #25269 (2024-12-04 22:29:14)

[Pallas] Pallas documentation cleanup
Minor docs cleanup.

- Move Pallas Design and Async Notes to a new "Design Notes" folder.
- Remove a duplicate example between quickstart and block specs.
- Some minor rewordings in quickstart.


Patch:
diff --git a/docs/pallas/async_note.md b/docs/pallas/design/async_note.md
similarity index 100%
rename from docs/pallas/async_note.md
rename to docs/pallas/design/async_note.md
diff --git a/docs/pallas/design.md b/docs/pallas/design/design.md
similarity index 99%
rename from docs/pallas/design.md
rename to docs/pallas/design/design.md
index f6fc8f592..17c7a6dbd 100644
--- a/docs/pallas/design.md
+++ b/docs/pallas/design/design.md
@@ -94,7 +94,7 @@ Pallas kernels via JAX transformations.
 
 <center>
 
-![Pallas lowering path](../_static/pallas/pallas_flow.png)
+![Pallas lowering path](../../_static/pallas/pallas_flow.png)
 Visualization of Pallas lowering paths
 
 </center>
@@ -413,10 +413,10 @@ verify the correctness of the Triton and Mosaic compilers.
 One could also imagine perturbing the `scan` ordering to simulate the
 parallel reads and writes that happen on GPU.
 
-### Examples
+### GPU Examples
 
-Note all the following examples are for GPU only. They will require some small
-changes to work on TPUs.
+Note all the following examples are for GPU only. They will require tweaks to
+the block sizes to work on TPUs.
 
 #### `add`
 
diff --git a/docs/pallas/design/index.rst b/docs/pallas/design/index.rst
new file mode 100644
index 000000000..d11a13d39
--- /dev/null
+++ b/docs/pallas/design/index.rst
@@ -0,0 +1,9 @@
+Pallas Design Notes
+===================
+
+.. toctree::
+   :caption: Design
+   :maxdepth: 2
+
+   design
+   async_note
diff --git a/docs/pallas/grid_blockspec.md b/docs/pallas/grid_blockspec.md
index cde200528..c1b2c2b95 100644
--- a/docs/pallas/grid_blockspec.md
+++ b/docs/pallas/grid_blockspec.md
@@ -44,39 +44,7 @@ For example, for invocation `(1, 2)`, `program_id(axis=0)` returns `1` and
 You can also use {func}`jax.experimental.pallas.num_programs` to get the
 grid size for a given axis.
 
-Here's an example kernel that uses a `grid` and `program_id`.
-
-```python
->>> import jax
->>> from jax.experimental import pallas as pl
-
->>> def iota_kernel(o_ref):
-...   i = pl.program_id(0)
-...   o_ref[i] = i
-
-```
-
-We now execute it using `pallas_call` with an additional `grid` argument.
-
-```python
->>> def iota(size: int):
-...   return pl.pallas_call(iota_kernel,
-...                         out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),
-...                         grid=(size,), interpret=True)()
->>> iota(8)
-Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)
-
-```
-
-On GPUs, each program is executed in parallel on separate thread blocks.
-Thus, we need to think about race conditions on writes to HBM.
-A reasonable approach is to write our kernels in such a way that different
-programs write to disjoint places in HBM to avoid these parallel writes.
-
-On TPUs, programs are executed in a combination of parallel and sequential
-(depending on the architecture) so there are slightly different considerations.
-
-See {ref}`pallas_tpu_noteworthy_properties`.
+See {ref}`grids_by_example` for a simple kernel that uses this API.
 
 (pallas_blockspec)=
 
@@ -131,6 +99,8 @@ shape `x_shape` are computed as in the function `slice_for_invocation`
 below:
 
 ```python
+>>> import jax
+>>> from jax.experimental import pallas as pl
 >>> def slices_for_invocation(x_shape: tuple[int, ...],
 ...                           x_spec: pl.BlockSpec,
 ...                           grid: tuple[int, ...],
diff --git a/docs/pallas/index.rst b/docs/pallas/index.rst
index 5969349c9..b2e2fca6c 100644
--- a/docs/pallas/index.rst
+++ b/docs/pallas/index.rst
@@ -22,7 +22,6 @@ See also the :class:`jax.experimental.pallas` module API documentation.
    :maxdepth: 2
 
    quickstart
-   design
    grid_blockspec
 
 
@@ -34,9 +33,9 @@ See also the :class:`jax.experimental.pallas` module API documentation.
 
 .. toctree::
    :caption: Design Notes
-   :maxdepth: 1
+   :maxdepth: 2
 
-   async_note
+   design/index
 
 .. toctree::
    :caption: Other
diff --git a/docs/pallas/quickstart.ipynb b/docs/pallas/quickstart.ipynb
index af34d1674..11dd2108e 100644
--- a/docs/pallas/quickstart.ipynb
+++ b/docs/pallas/quickstart.ipynb
@@ -72,8 +72,9 @@
     "\n",
     "Let's dissect this function a bit. Unlike most JAX functions you've probably written,\n",
     "it does not take in `jax.Array`s as inputs and doesn't return any values.\n",
-    "Instead, it takes in *`Ref`* objects as inputs. Note that we also don't have any outputs\n",
-    "but we are given an `o_ref`, which corresponds to the desired output.\n",
+    "Instead, it takes in *`Ref`* objects as inputs, which represent mutable buffers in memory.\n",
+    "Note that we also don't have any outputs but we are given an `o_ref`, which corresponds\n",
+    "to the desired output.\n",
     "\n",
     "**Reading from `Ref`s**\n",
     "\n",
@@ -150,7 +151,8 @@
     "**What's actually happening here?**\n",
     "\n",
     "Thus far we've described how to think about Pallas kernels but what we've actually\n",
-    "accomplished is we're writing a function that's executed very close to the compute units.\n",
+    "accomplished is we're writing a function that's executed very close to the compute units\n",
+    "since values are loaded into the innermost (fastest) portion of the memory hierarchy.\n",
     "\n",
     "On GPU, `x_ref` corresponds to a value in high-bandwidth memory (HBM) and when\n",
     "we do `x_ref[...]` we are copying the value from HBM into static RAM (SRAM)\n",
@@ -195,6 +197,8 @@
     "live in high-bandwidth memory (HBM, also known as DRAM) and expressing computations\n",
     "that operate on \"blocks\" of those arrays that can fit in SRAM.\n",
     "\n",
+    "(grids_by_example)=\n",
+    "\n",
     "### Grids by example\n",
     "\n",
     "To automatically \"carve\" up the inputs and outputs, you provide a `grid` and\n",
@@ -240,7 +244,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We now execute it using `pallas_call` with an additional `grid` argument."
+    "We now execute it using `pallas_call` with an additional `grid` argument.\n",
+    "On GPUs, we can call the kernel directly like so:"
    ]
   },
   {
@@ -260,6 +265,7 @@
     }
    ],
    "source": [
+    "# GPU version\n",
     "def iota(size: int):\n",
     "  return pl.pallas_call(iota_kernel,\n",
     "                        out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),\n",
@@ -272,16 +278,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "On GPUs, each program is executed in parallel on separate threads.\n",
-    "Thus, we need to think about race conditions on writes to HBM.\n",
-    "A reasonable approach is to write our kernels in such a way that different\n",
-    "programs write to disjoint places in HBM to avoid these parallel writes.\n",
-    "On the other hand, parallelizing the computation is how we can execute\n",
-    "operations like matrix multiplications really quickly.\n",
-    "\n",
-    "On TPUs, programs are executed in a combination of parallel and sequential\n",
-    "(depending on the architecture) so there are slightly different considerations.\n",
-    "\n",
+    "TPUs distinguish between vector and scalar memory spaces and in this case the\n",
+    "output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is\n",
+    "a scalar. For more details read {ref}`tpu_and_its_memory_spaces`.\n",
     "To call the above kernel on TPU, run:"
    ]
   },
@@ -292,6 +291,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# TPU version\n",
     "from jax.experimental.pallas import tpu as pltpu\n",
     "\n",
     "def iota(size: int):\n",
@@ -307,11 +307,22 @@
    "id": "68f97b4e",
    "metadata": {},
    "source": [
-    "TPUs distinguish between vector and scalar memory spaces and in this case the\n",
-    "output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is\n",
-    "a scalar. For more details read {ref}`pallas_tpu_pipelining`.\n",
+    "### Grid semantics\n",
+    "\n",
+    "On GPUs, each program is executed in parallel on separate threads.\n",
+    "Thus, we need to think about race conditions on writes to HBM.\n",
+    "A reasonable approach is to write our kernels in such a way that different\n",
+    "programs write to disjoint locations in HBM to avoid these parallel writes.\n",
+    "On the other hand, parallelizing the computation is how we can execute\n",
+    "operations like matrix multiplications really quickly.\n",
+    "\n",
+    "In contrast, TPUs operate like a very wide SIMD machine.\n",
+    "Some TPU models contain multiple cores, but in many cases a TPU can be\n",
+    "treated as a single-threaded processor. The grid on a TPU can be\n",
+    "specified in a combination of parallel and sequential dimensions, where sequential\n",
+    "dimensions are guaranteed to run serially.\n",
     "\n",
-    "You can read more details at {ref}`pallas_grid`."
+    "You can read more details at {ref}`pallas_grid` and {ref}`pallas_tpu_noteworthy_properties`."
    ]
   },
   {
diff --git a/docs/pallas/quickstart.md b/docs/pallas/quickstart.md
index e11868f5f..fff1dcb73 100644
--- a/docs/pallas/quickstart.md
+++ b/docs/pallas/quickstart.md
@@ -53,8 +53,9 @@ def add_vectors_kernel(x_ref, y_ref, o_ref):
 
 Let's dissect this function a bit. Unlike most JAX functions you've probably written,
 it does not take in `jax.Array`s as inputs and doesn't return any values.
-Instead, it takes in *`Ref`* objects as inputs. Note that we also don't have any outputs
-but we are given an `o_ref`, which corresponds to the desired output.
+Instead, it takes in *`Ref`* objects as inputs, which represent mutable buffers in memory.
+Note that we also don't have any outputs but we are given an `o_ref`, which corresponds
+to the desired output.
 
 **Reading from `Ref`s**
 
@@ -101,7 +102,8 @@ thereof).
 **What's actually happening here?**
 
 Thus far we've described how to think about Pallas kernels but what we've actually
-accomplished is we're writing a function that's executed very close to the compute units.
+accomplished is we're writing a function that's executed very close to the compute units
+since values are loaded into the innermost (fastest) portion of the memory hierarchy.
 
 On GPU, `x_ref` corresponds to a value in high-bandwidth memory (HBM) and when
 we do `x_ref[...]` we are copying the value from HBM into static RAM (SRAM)
@@ -134,6 +136,8 @@ Part of writing Pallas kernels is thinking about how to take big arrays that
 live in high-bandwidth memory (HBM, also known as DRAM) and expressing computations
 that operate on "blocks" of those arrays that can fit in SRAM.
 
+(grids_by_example)=
+
 ### Grids by example
 
 To automatically "carve" up the inputs and outputs, you provide a `grid` and
@@ -169,8 +173,10 @@ def iota_kernel(o_ref):
 ```
 
 We now execute it using `pallas_call` with an additional `grid` argument.
+On GPUs, we can call the kernel directly like so:
 
 ```{code-cell} ipython3
+# GPU version
 def iota(size: int):
   return pl.pallas_call(iota_kernel,
                         out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),
@@ -178,19 +184,13 @@ def iota(size: int):
 iota(8)
 ```
 
-On GPUs, each program is executed in parallel on separate threads.
-Thus, we need to think about race conditions on writes to HBM.
-A reasonable approach is to write our kernels in such a way that different
-programs write to disjoint places in HBM to avoid these parallel writes.
-On the other hand, parallelizing the computation is how we can execute
-operations like matrix multiplications really quickly.
-
-On TPUs, programs are executed in a combination of parallel and sequential
-(depending on the architecture) so there are slightly different considerations.
-
+TPUs distinguish between vector and scalar memory spaces and in this case the
+output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is
+a scalar. For more details read {ref}`tpu_and_its_memory_spaces`.
 To call the above kernel on TPU, run:
 
 ```{code-cell} ipython3
+# TPU version
 from jax.experimental.pallas import tpu as pltpu
 
 def iota(size: int):
@@ -201,11 +201,22 @@ def iota(size: int):
 iota(8)
 ```
 
-TPUs distinguish between vector and scalar memory spaces and in this case the
-output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is
-a scalar. For more details read {ref}`pallas_tpu_pipelining`.
+### Grid semantics
+
+On GPUs, each program is executed in parallel on separate threads.
+Thus, we need to think about race conditions on writes to HBM.
+A reasonable approach is to write our kernels in such a way that different
+programs write to disjoint locations in HBM to avoid these parallel writes.
+On the other hand, parallelizing the computation is how we can execute
+operations like matrix multiplications really quickly.
+
+In contrast, TPUs operate like a very wide SIMD machine.
+Some TPU models contain multiple cores, but in many cases a TPU can be
+treated as a single-threaded processor. The grid on a TPU can be
+specified in a combination of parallel and sequential dimensions, where sequential
+dimensions are guaranteed to run serially.
 
-You can read more details at {ref}`pallas_grid`.
+You can read more details at {ref}`pallas_grid` and {ref}`pallas_tpu_noteworthy_properties`.
 
 +++
 
diff --git a/docs/pallas/tpu/pipelining.ipynb b/docs/pallas/tpu/pipelining.ipynb
index 9774e08dc..10de58710 100644
--- a/docs/pallas/tpu/pipelining.ipynb
+++ b/docs/pallas/tpu/pipelining.ipynb
@@ -48,12 +48,20 @@
   },
   {
    "cell_type": "markdown",
+   "id": "0e212a5e",
    "metadata": {
     "id": "TWKESTKAlyjT"
    },
    "source": [
-    "## TPU and its memory spaces\n",
+    "(tpu_and_its_memory_spaces)=\n",
     "\n",
+    "## TPU and its memory spaces"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "A TPU and its TensorCore consist of memory spaces (where arrays can reside),\n",
     "registers (which temporarily store scalar and array values) and compute units\n",
     "(that do computation with values in registers).\n",
diff --git a/docs/pallas/tpu/pipelining.md b/docs/pallas/tpu/pipelining.md
index 218654301..df570cf08 100644
--- a/docs/pallas/tpu/pipelining.md
+++ b/docs/pallas/tpu/pipelining.md
@@ -38,8 +38,12 @@ import numpy as np
 
 +++ {"id": "TWKESTKAlyjT"}
 
+(tpu_and_its_memory_spaces)=
+
 ## TPU and its memory spaces
 
++++
+
 A TPU and its TensorCore consist of memory spaces (where arrays can reside),
 registers (which temporarily store scalar and array values) and compute units
 (that do computation with values in registers).
diff --git a/jax/_src/config.py b/jax/_src/config.py
index b0214a8c8..eea686fa0 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1079,9 +1079,8 @@ else:
     def __init__(self, name):
       self._name = name
 
-    @property
     def value(self):
-      return self.get_local()
+      return getattr(jax_jit.thread_local_state().extra_jit_context, self._name)
 
     def get_local(self):
       return getattr(jax_jit.thread_local_state().extra_jit_context, self._name)
@@ -1089,11 +1088,6 @@ else:
     def set_local(self, value):
       update_thread_local_jit_state(**{self._name: value})
 
-    def swap_local(self, new_value):
-      prev_value = self.value
-      self.set_local(new_value)
-      return prev_value
-
   trace_state = JitConfig('trace_state')
   axis_env_state = JitConfig('axis_env_state')
   mesh_context_manager = JitConfig('mesh_context_manager')
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 30893ce99..122b7bcf5 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1605,10 +1605,10 @@ def get_sharding(sharding, ndim):
     assert len(sharding.spec) == ndim
     return sharding
 
-  context_mesh = mesh_lib.get_abstract_mesh()
+  context_mesh = mesh_lib.abstract_mesh_context.mesh
   # TODO(yashkatariya): Error out and ask users to set the context mesh in their
   # code.
-  if not context_mesh:
+  if context_mesh is None:
     return None
   assert sharding is None
   return NamedSharding(context_mesh, P(*[None] * ndim))
diff --git a/jax/_src/deprecations.py b/jax/_src/deprecations.py
index 33a813340..c7a956068 100644
--- a/jax/_src/deprecations.py
+++ b/jax/_src/deprecations.py
@@ -125,6 +125,7 @@ register('jax-aval-named-shape')
 register('jax-dlpack-import-legacy')
 register("jax-numpy-astype-complex-to-real")
 register("jax-numpy-array-none")
+register('jax-numpy-reshape-newshape')
 register('jax-numpy-clip-args')
 register('jax-numpy-linalg-matrix_rank-tol')
 register('jax-numpy-linalg-pinv-rcond')
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index 87b931683..d7a87b6d6 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -454,6 +454,18 @@ class AbstractMesh:
   def local_mesh(self):
     _raise_value_error("local_mesh")
 
+  def __enter__(self):
+    abstract_mesh_context.stack.append(self)
+    abstract_mesh_context.mesh = self
+    jax_config.abstract_mesh_context_manager.set_local(abstract_mesh_context.mesh)
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    abstract_mesh_context.stack.pop()
+    abstract_mesh_context.mesh = abstract_mesh_context.stack[-1]
+    jax_config.abstract_mesh_context_manager.set_local(abstract_mesh_context.mesh)
+    return False
+
   @staticmethod
   def _extremely_unsafe_enter_tracing_context(mesh: AbstractMesh):
     jax_config.abstract_mesh_context_manager.set_local(mesh)
@@ -466,32 +478,37 @@ def _raise_value_error(name):
   raise ValueError(f"AbstractMesh does not implement {name}")
 
 
-@contextlib.contextmanager
-def set_abstract_mesh(mesh: AbstractMesh):
-  prev_val = jax_config.abstract_mesh_context_manager.swap_local(mesh)
-  try:
-    yield
-  finally:
-    jax_config.abstract_mesh_context_manager.set_local(prev_val)
+class AbstractMeshContext(threading.local):
+  def __init__(self):
+    self.stack = [None]
+    self.mesh = self.stack[-1]
 
-def get_abstract_mesh():
-  return jax_config.abstract_mesh_context_manager.value
+abstract_mesh_context = AbstractMeshContext()
 
 
 @contextlib.contextmanager
-def set_concrete_mesh(mesh: Mesh):
-  prev_val = jax_config.device_context.swap_local(mesh)
-  try:
+def set_mesh(mesh: Mesh):
+  with (mesh.abstract_mesh, jax_config.sharding_in_types(True),
+        enter_device_context(mesh)):
     yield
-  finally:
-    jax_config.device_context.set_local(prev_val)
 
-def get_concrete_mesh():
-  return jax_config.device_context.value
+
+class DeviceContext(threading.local):
+  def __init__(self):
+    self.stack = [None]
+    self.concrete_mesh = self.stack[-1]
+
+device_context = DeviceContext()
 
 
 @contextlib.contextmanager
-def set_mesh(mesh: Mesh):
-  with (set_abstract_mesh(mesh.abstract_mesh),
-        jax_config.sharding_in_types(True), set_concrete_mesh(mesh)):
+def enter_device_context(mesh: Mesh):
+  device_context.stack.append(mesh)
+  device_context.concrete_mesh = mesh
+  jax_config.device_context.set_local(device_context.concrete_mesh)
+  try:
     yield
+  finally:
+    device_context.stack.pop()
+    device_context.concrete_mesh = device_context.stack[-1]
+    jax_config.device_context.set_local(device_context.concrete_mesh)
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 6c632c6fc..1f8378d84 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -16,6 +16,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Callable, Sequence, Iterable
+import contextlib
 import dataclasses
 from functools import partial
 import inspect
@@ -186,7 +187,7 @@ def _python_pjit_helper(fun, jit_info, *args, **kwargs):
   try:
     # TODO(yashkatariya): Maybe thread this into pjit params like resource_env
     # and set the context manager down the stack?
-    with mesh_lib.set_abstract_mesh(p.abstract_mesh):
+    with p.abstract_mesh:
       if (core.trace_state_clean() and
           not config.debug_key_reuse.value and
           not config.data_dependent_tracing_fallback.value):
@@ -644,9 +645,9 @@ def _infer_params_impl(
   attr_token = _attr_token(flat_fun, in_type)
 
   abstract_mesh = (
-      get_abstract_mesh_from_avals(in_type)
-      if not mesh_lib.get_abstract_mesh() else mesh_lib.get_abstract_mesh())
-  with mesh_lib.set_abstract_mesh(abstract_mesh):
+      get_abstract_mesh(in_type) if mesh_lib.abstract_mesh_context.mesh is None
+      else mesh_lib.abstract_mesh_context.mesh)
+  with abstract_mesh:
     jaxpr, consts, out_avals, attrs_tracked = _create_pjit_jaxpr(
         flat_fun, in_type, attr_token, dbg,
         HashableFunction(res_paths, closure=()),
@@ -693,9 +694,9 @@ def _infer_params_impl(
                     attrs_tracked, abstract_mesh), args_flat
 
 
-def get_abstract_mesh_from_avals(in_avals):
+def get_abstract_mesh(in_avals):
   if not config.sharding_in_types.value:
-    return None
+    return contextlib.nullcontext()
   m = None
   for a in in_avals:
     # TODO(yashkatariya): Remove this when mesh context can be set by the user.
@@ -1788,8 +1789,7 @@ def _pjit_lower(
     lowering_parameters: mlir.LoweringParameters,
     pgle_profiler: profiler.PGLEProfiler | None):
   if config.sharding_in_types.value:
-    cur_mesh = mesh_lib.get_concrete_mesh()
-    mesh = cur_mesh if isinstance(cur_mesh, mesh_lib.Mesh) else None
+    mesh = mesh_lib.device_context.concrete_mesh
     api_name = 'jit'
   else:
     mesh, api_name = ((resource_env.physical_mesh, 'pjit')
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index db26813de..b6f3b63d3 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -30,6 +30,7 @@ executable protocols described above.
 """
 from __future__ import annotations
 
+import contextlib
 import functools
 from collections.abc import Sequence
 from dataclasses import dataclass
@@ -43,7 +44,6 @@ from jax._src import source_info_util
 from jax._src import traceback_util
 from jax._src import tree_util
 from jax._src import util
-from jax._src import mesh as mesh_lib
 from jax._src.sharding_impls import UnspecifiedValue, AUTO
 from jax._src.layout import Layout
 from jax._src.interpreters import mlir
@@ -717,7 +717,7 @@ class Traced(Stage):
                "_args_flat", "_arg_names", "_num_consts"]
 
   def __init__(self, jaxpr: core.ClosedJaxpr, args_info, fun_name, out_tree,
-               lower_callable, abstract_mesh=None,
+               lower_callable, abstract_mesh=contextlib.nullcontext(),
                args_flat=None, arg_names=None, num_consts: int = 0):
     self.jaxpr = jaxpr
     self.args_info = args_info
@@ -747,7 +747,7 @@ class Traced(Stage):
     try:
       # TODO(yashkatariya): Maybe thread this into pjit params like resource_env
       # and set the context manager down the stack?
-      with mesh_lib.set_abstract_mesh(self._abstract_mesh):
+      with self._abstract_mesh:
         lowering = new_callable()
     except pxla.DeviceAssignmentMismatchError as e:
       fails, = e.args
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index 1c529b893..07f631f6e 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -46,7 +46,7 @@ from jax._src import source_info_util
 from jax._src import traceback_util
 from jax._src import util
 from jax._src.core import Tracer
-from jax._src.mesh import AbstractMesh, Mesh, AxisTypes, set_abstract_mesh
+from jax._src.mesh import AbstractMesh, Mesh, AxisTypes
 from jax._src.api import _shared_code_pmap, _prepare_pmap
 from jax._src.lax import (lax, parallel as lax_parallel, slicing,
                           windowed_reductions, convolution, fft, linalg,
@@ -484,7 +484,7 @@ def _shard_map_staging(
   in_avals = [t.aval for t in in_tracers]
   in_avals_ = map(partial(_shard_aval, mesh), in_names, in_avals)
   with (core.extend_axis_env_nd(list(mesh.shape.items())),
-        set_abstract_mesh(pjit.get_abstract_mesh_from_avals(in_avals_))):
+        pjit.get_abstract_mesh(in_avals_)):
     jaxpr, out_avals_, consts, () = pe.trace_to_jaxpr_dynamic(f, in_avals_)
   _check_names(out_names_thunk(), out_avals_)
   if check_rep:


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 7e66c41cc..52261ef02 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5237,7 +5237,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     def g(x, y):
       self.assertTrue(x.sharding.mesh._are_all_axes_collective)
       self.assertTrue(y.sharding.mesh._are_all_axes_collective)
-      self.assertTrue(mesh_lib.get_abstract_mesh()._are_all_axes_collective)
+      self.assertTrue(mesh_lib.abstract_mesh_context.mesh._are_all_axes_collective)
       return x * y
 
     @jax.jit
@@ -5262,7 +5262,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     def g(x, y):
       self.assertTrue(x.sharding.mesh._are_all_axes_collective)
       self.assertTrue(y.sharding.mesh._are_all_axes_collective)
-      self.assertTrue(mesh_lib.get_abstract_mesh()._are_all_axes_collective)
+      self.assertTrue(mesh_lib.abstract_mesh_context.mesh._are_all_axes_collective)
       allgatherd_y = jax.lax.all_gather(y, axis_name='x', axis=1, tiled=True)
       z = x @ allgatherd_y
       return jax.lax.psum(z, axis_name='y')