Repository: jax-ml/jax
Pull request #25290 (2024-12-05 17:26:10)

jax.numpy: require boolean dtype for where argument
This requirement is documented, but not currently validated.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b6d0f97f4..ce8b04043 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,9 +53,6 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
          use `uses_global_constants`.
       * the `lowering_platforms` kwarg for {func}`jax.export.export`: use
         `platforms` instead.
-  * The kwargs `symbolic_scope` and `symbolic_constraints` from
-    {func}`jax.export.symbolic_args_specs` have been removed. They were
-    deprecated in June 2024. Use `scope` and `constraints` instead.
   * Hashing of tracers, which has been deprecated since version 0.4.30, now
     results in a `TypeError`.
   * Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and
@@ -70,11 +67,6 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     return NaN for negative integer inputs, to match the behavior of SciPy from
     https://github.com/scipy/scipy/pull/21827.
   * `jax.clear_backends` was removed after being deprecated in v0.4.26.
-  * We removed the custom call "__gpu$xla.gpu.triton" from the list of custom
-    call that we guarantee export stability. This is because this custom call
-    relies on Triton IR, which is not guaranteed to be stable. If you need
-    to export code that uses this custom call, you can use the `disabled_checks`
-    parameter. See more details in the [documentation](https://jax.readthedocs.io/en/latest/export/export.html#compatibility-guarantees-for-custom-calls).
 
 * New Features
   * {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
@@ -87,7 +79,6 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.lax.linalg.eig` and the related `jax.numpy` functions
     ({func}`jax.numpy.linalg.eig` and {func}`jax.numpy.linalg.eigvals`) are now
     supported on GPU. See {jax-issue}`#24663` for more details.
-  * Added two new configuration flags, `jax_exec_time_optimization_effort` and `jax_memory_fitting_effort`, to control the amount of effort the compiler spends minimizing execution time and memory usage, respectively.  Valid values are between -1.0 and 1.0, default is 0.0.
 
 * Bug fixes
   * Fixed a bug where the GPU implementations of LU and QR decomposition would
diff --git a/jax/_src/ad_checkpoint.py b/jax/_src/ad_checkpoint.py
index 93376c7bd..fc135ac8f 100644
--- a/jax/_src/ad_checkpoint.py
+++ b/jax/_src/ad_checkpoint.py
@@ -652,7 +652,7 @@ def remat_transpose(out_cts, *in_primals, jaxpr, **params):
             for x in in_primals]
   assert next(in_cts_nz_, None) is next(in_zeros_, None) is None
   return in_cts
-ad.primitive_transposes[remat_p] = remat_transpose
+ad.reducing_transposes[remat_p] = remat_transpose
 
 # TODO(mattjj): move this to ad.py
 def transpose_jaxpr(jaxpr: core.ClosedJaxpr, in_linear: bool | Sequence[bool],
diff --git a/jax/_src/cache_key.py b/jax/_src/cache_key.py
index 2ec645cee..6e7a42148 100644
--- a/jax/_src/cache_key.py
+++ b/jax/_src/cache_key.py
@@ -238,38 +238,6 @@ def _hash_accelerator_config(hash_obj, accelerators: np.ndarray, backend):
     _hash_devices(hash_obj, accelerators)
     _hash_platform(hash_obj, backend)
 
-# LINT.IfChange(xla_flags)
-xla_flags_to_exclude_from_cache_key = [
-    "--xla_dump_compress_protos",
-    "--xla_dump_module_metadata",
-    "--xla_dump_max_hlo_modules",
-    "--xla_dump_include_timestamp",
-    "--xla_dump_hlo_pass_re",
-    "--xla_dump_hlo_module_re",
-    "--xla_dump_hlo_snapshots",
-    "--xla_dump_fusion_visualization",
-    "--xla_dump_hlo_as_url",
-    "--xla_dump_hlo_as_proto",
-    "--xla_dump_hlo_as_text",
-    "--xla_dump_hlo_as_long_text",
-    "--xla_dump_hlo_as_html",
-    "--xla_dump_hlo_as_dot",
-    "--xla_dump_to",
-    "--xla_force_host_platform_device_count",
-    "--xla_dump_disable_metadata",
-    "--xla_dump_hlo_pipeline_re",
-    "--xla_tpu_sdc_checker_streamz_metric",
-    "--xla_tpu_sdc_checker_enable_sdc_event_callbacks",
-    "--xla_tpu_sdc_checker_enable_coresweep_ng_callbacks",
-    "--xla_tpu_sdc_checker_no_logging_if_callbacks_are_present",
-    "--xla_gpu_cuda_data_dir",
-    "--xla_gpu_experimental_autotune_cache_mode",
-]
-
-env_override_flags_to_exclude_from_cache_key = {
-    x.strip("-") for x in xla_flags_to_exclude_from_cache_key
-}
-# LINT.ThenChange(:debug_options)
 
 def _hash_serialized_compile_options(hash_obj, compile_options_obj,
                                      strip_device_assignment=False):
@@ -316,11 +284,6 @@ def _hash_serialized_compile_options(hash_obj, compile_options_obj,
   debug_options.xla_gpu_cuda_data_dir = ""
   # LINT.ThenChange(:xla_flags)
 
-  compile_options_copy.env_option_overrides = [
-      flag_value
-      for flag_value in compile_options_copy.env_option_overrides
-      if flag_value[0] not in env_override_flags_to_exclude_from_cache_key
-  ]
   if strip_device_assignment and compile_options_copy.device_assignment:
     replica_count = compile_options_copy.device_assignment.replica_count()
     computation_count = compile_options_copy.device_assignment.computation_count()
@@ -338,6 +301,34 @@ def _hash_platform(hash_obj, backend):
 
 
 def _hash_xla_flags(hash_obj, extra_flag_prefixes: list[str]):
+  # LINT.IfChange(xla_flags)
+  xla_flags_to_exclude_from_cache_key = [
+      "--xla_dump_compress_protos",
+      "--xla_dump_module_metadata",
+      "--xla_dump_max_hlo_modules",
+      "--xla_dump_include_timestamp",
+      "--xla_dump_hlo_pass_re",
+      "--xla_dump_hlo_module_re",
+      "--xla_dump_hlo_snapshots",
+      "--xla_dump_fusion_visualization",
+      "--xla_dump_hlo_as_url",
+      "--xla_dump_hlo_as_proto",
+      "--xla_dump_hlo_as_text",
+      "--xla_dump_hlo_as_long_text",
+      "--xla_dump_hlo_as_html",
+      "--xla_dump_hlo_as_dot",
+      "--xla_dump_to",
+      "--xla_force_host_platform_device_count",
+      "--xla_dump_disable_metadata",
+      "--xla_dump_hlo_pipeline_re",
+      "--xla_tpu_sdc_checker_streamz_metric",
+      "--xla_tpu_sdc_checker_enable_sdc_event_callbacks",
+      "--xla_tpu_sdc_checker_enable_coresweep_ng_callbacks",
+      "--xla_tpu_sdc_checker_no_logging_if_callbacks_are_present",
+      "--xla_gpu_cuda_data_dir",
+  ]
+  # LINT.ThenChange(:debug_options)
+
   xla_flags = []
 
   xla_flags_env_var = os.getenv("XLA_FLAGS")
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index 2d032069e..d40ded556 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -36,7 +36,6 @@ from jax._src import traceback_util
 from jax._src.interpreters import mlir
 from jax._src.lib import version as jaxlib_version
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 import numpy as np
 
@@ -192,10 +191,6 @@ def get_compile_options(
     assert device_assignment.computation_count() == num_partitions
     compile_options.device_assignment = device_assignment
 
-  if xla_extension_version >= 294:
-    build_options.exec_time_optimization_effort = config.exec_time_optimization_effort.value
-    build_options.memory_fitting_effort = config.memory_fitting_effort.value
-
   if env_options_overrides is not None:
     # Some overrides are passed directly on build_options.
     overrides_on_build_options = [
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 34a1e152c..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')
@@ -2006,15 +2000,3 @@ gpu_use_magma = enum_state(
         'to use this feature.'
     ),
 )
-
-exec_time_optimization_effort = float_state(
-    name='jax_exec_time_optimization_effort',
-    default=0.0,
-    help='Effort for minimizing execution time (higher means more effort), valid range [-1.0, 1.0].'
-)
-
-memory_fitting_effort = float_state(
-    name='jax_memory_fitting_effort',
-    default=0.0,
-    help='Effort for minimizing memory usage (higher means more effort), valid range [-1.0, 1.0].'
-)
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..778a084e8 100644
--- a/jax/_src/deprecations.py
+++ b/jax/_src/deprecations.py
@@ -125,9 +125,11 @@ 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')
 register('jax-numpy-quantile-interpolation')
+register('jax-numpy-reduction-non-boolean-where')
 register('jax-numpy-trimzeros-not-1d-array')
 register('pallas-gpu-triton')
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
index e3508639f..ad2c7fdac 100644
--- a/jax/_src/export/_export.py
+++ b/jax/_src/export/_export.py
@@ -1005,8 +1005,7 @@ _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE = {
     *_CPU_FFI_KERNELS,
     "Sharding", "SPMDFullToShardShape", "SPMDShardToFullShape",
     "cu_threefry2x32", "cu_threefry2x32_ffi",
-    # Triton IR does not guarantee stability.
-    # "__gpu$xla.gpu.triton",
+    "__gpu$xla.gpu.triton",  # Pallas call on GPU
     # cholesky on CPU
     "lapack_spotrf", "lapack_dpotrf", "lapack_cpotrf", "lapack_zpotrf",
     # eigh on TPU
diff --git a/jax/_src/export/shape_poly.py b/jax/_src/export/shape_poly.py
index cb9a99564..15f99533d 100644
--- a/jax/_src/export/shape_poly.py
+++ b/jax/_src/export/shape_poly.py
@@ -1198,6 +1198,12 @@ def is_symbolic_dim(p: DimSize) -> bool:
   """
   return isinstance(p, _DimExpr)
 
+def is_poly_dim(p: DimSize) -> bool:
+  # TODO: deprecated January 2024, remove June 2024.
+  warnings.warn("is_poly_dim is deprecated, use export.is_symbolic_dim",
+                DeprecationWarning, stacklevel=2)
+  return is_symbolic_dim(p)
+
 dtypes.python_scalar_dtypes[_DimExpr] = dtypes.python_scalar_dtypes[int]
 
 def _einsum_contract_path(*operands, **kwargs):
@@ -1407,6 +1413,8 @@ def symbolic_args_specs(
     shapes_specs,  # prefix pytree of strings
     constraints: Sequence[str] = (),
     scope: SymbolicScope | None = None,
+    symbolic_constraints: Sequence[str] = (),  # DEPRECATED on 6/14/24
+    symbolic_scope: SymbolicScope | None = None,  # DEPRECATED on 6/14/24
 ):
   """Constructs a pytree of jax.ShapeDtypeSpec arguments specs for `export`.
 
@@ -1427,10 +1435,25 @@ def symbolic_args_specs(
       arguments](https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees).
     constraints: as for :func:`jax.export.symbolic_shape`.
     scope: as for :func:`jax.export.symbolic_shape`.
+    symbolic_constraints: DEPRECATED, use `constraints`.
+    symbolic_scope: DEPRECATED, use `scope`.
 
   Returns: a pytree of jax.ShapeDTypeStruct matching the `args` with the shapes
     replaced with symbolic dimensions as specified by `shapes_specs`.
   """
+  if symbolic_constraints:
+    warnings.warn("symbolic_constraints is deprecated, use constraints",
+                  DeprecationWarning, stacklevel=2)
+    if constraints:
+      raise ValueError("Cannot use both symbolic_constraints and constraints")
+    constraints = symbolic_constraints
+  if symbolic_scope is not None:
+    warnings.warn("symbolic_scope is deprecated, use scope",
+                  DeprecationWarning, stacklevel=2)
+    if scope is not None:
+      raise ValueError("Cannot use both symbolic_scope and scope")
+    scope = symbolic_scope
+
   polymorphic_shapes = shapes_specs
   args_flat, args_tree = tree_util.tree_flatten(args)
 
diff --git a/jax/_src/interpreters/ad.py b/jax/_src/interpreters/ad.py
index c5e78321f..804df185d 100644
--- a/jax/_src/interpreters/ad.py
+++ b/jax/_src/interpreters/ad.py
@@ -277,6 +277,9 @@ def backward_pass(jaxpr: core.Jaxpr, transform_stack,
           call_jaxpr = params.pop('call_jaxpr')
           cts_out = get_primitive_transpose(eqn.primitive)(
               params, call_jaxpr, invals, cts_in, cts_in_avals)
+        elif eqn.primitive in reducing_transposes:
+          cts_out = reducing_transposes[eqn.primitive](
+              cts_in, *invals, **eqn.params)
         else:
           cts_out = get_primitive_transpose(eqn.primitive)(
               cts_in, *invals, **eqn.params)
@@ -583,6 +586,8 @@ class LinearizeTracer(Tracer):
 
 primitive_jvps : dict[core.Primitive, Callable] = {}
 primitive_transposes: dict[core.Primitive, Callable] = {}
+# transpose rules that internally perform reductions over the given named axes
+reducing_transposes: dict[core.Primitive, Callable] = {}
 primitive_linearizations : dict[core.Primitive, Callable]  = {}
 
 def deflinear(primitive, transpose_rule):
@@ -866,6 +871,3 @@ class CustomVJPException(Exception):
            "closed-over value into the custom_vjp function as an argument, and "
            "adapting the custom_vjp fwd and bwd rules.")
     super().__init__(msg)
-
-# TODO(mattjj): remove this vestigial dict
-reducing_transposes: dict[core.Primitive, Callable] = {}
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index 418240a4a..7a5e9e9de 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -780,7 +780,7 @@ cond_p.multiple_results = True
 cond_p.def_impl(partial(dispatch.apply_primitive, cond_p))
 cond_p.def_effectful_abstract_eval(_cond_abstract_eval)
 ad.primitive_jvps[cond_p] = _cond_jvp
-ad.primitive_transposes[cond_p] = _cond_transpose
+ad.reducing_transposes[cond_p] = _cond_transpose
 pe.custom_partial_eval_rules[cond_p] = _cond_partial_eval
 batching.fancy_primitive_batchers[cond_p] = _cond_batching_rule
 xla.register_initial_style_primitive(cond_p)
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index f62ce2434..76132ccdc 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -1228,7 +1228,7 @@ scan_p.multiple_results = True
 scan_p.def_impl(partial(dispatch.apply_primitive, scan_p))
 scan_p.def_effectful_abstract_eval(_scan_abstract_eval)
 ad.primitive_jvps[scan_p] = _scan_jvp
-ad.primitive_transposes[scan_p] = _scan_transpose
+ad.reducing_transposes[scan_p] = _scan_transpose
 pe.custom_partial_eval_rules[scan_p] = _scan_partial_eval
 xla.register_initial_style_primitive(scan_p)
 mlir.register_lowering(scan_p,
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 9c8afe4f9..ad08e1335 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -923,9 +923,8 @@ class DotAlgorithmPreset(enum.Enum):
       case _:
         return np.float32
 
-  def supported_output_types(
-      self, lhs_dtype: DTypeLike, rhs_dtype: DTypeLike
-  ) -> tuple[DTypeLike, ...] | None:
+  @property
+  def supported_output_types(self) -> tuple[DTypeLike, ...] | None:
     match self:
       case (
           DotAlgorithmPreset.ANY_F8_ANY_F8_F32
@@ -942,17 +941,7 @@ class DotAlgorithmPreset(enum.Enum):
             dtypes.float8_e4m3b11fnuz,
         )
       case DotAlgorithmPreset.F16_F16_F32:
-        # F16 output is only supported with F16 inputs.
-        if dtypes.promote_types(lhs_dtype, rhs_dtype) == np.float16:
-          return (np.float32, np.float16)
-        else:
-          return (np.float32,)
-      case DotAlgorithmPreset.BF16_BF16_F32:
-        # BF16 output is only supported with BF16 inputs.
-        if dtypes.promote_types(lhs_dtype, rhs_dtype) == dtypes.bfloat16:
-          return (np.float32, dtypes.bfloat16)
-        else:
-          return (np.float32,)
+        return (np.float32, np.float16)
       case _:
         accumulation_type = self.accumulation_type
         return None if accumulation_type is None else (accumulation_type,)
@@ -3724,6 +3713,13 @@ def get_algorithm_compute_types(
         algorithm.accumulation_type,
     )
 
+  supported_output_types = algorithm.supported_output_types
+
+  if algorithm == DotAlgorithmPreset.BF16_BF16_F32:
+    # If dtype is anything other than float32, it will be cast to bfloat16.
+    if np.dtype(lhs_dtype) != np.float32:
+      supported_output_types = (np.float32, dtypes.bfloat16)
+
   def maybe_convert_dtype(input_dtype, target_dtypes):
     if target_dtypes is None:
       return input_dtype
@@ -3731,12 +3727,11 @@ def get_algorithm_compute_types(
       return input_dtype
     return target_dtypes[0]
 
-  lhs_dtype = maybe_convert_dtype(lhs_dtype, algorithm.supported_lhs_types)
-  rhs_dtype = maybe_convert_dtype(rhs_dtype, algorithm.supported_rhs_types)
-  out_type = maybe_convert_dtype(
-      out_dtype, algorithm.supported_output_types(lhs_dtype, rhs_dtype)
+  return (
+      maybe_convert_dtype(lhs_dtype, algorithm.supported_lhs_types),
+      maybe_convert_dtype(rhs_dtype, algorithm.supported_rhs_types),
+      maybe_convert_dtype(out_dtype, supported_output_types),
   )
-  return lhs_dtype, rhs_dtype, out_type
 
 
 def _dot_general_lower(ctx, lhs, rhs, *, dimension_numbers,
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index c7b8f6920..d7a87b6d6 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -455,10 +455,16 @@ class AbstractMesh:
     _raise_value_error("local_mesh")
 
   def __enter__(self):
-    _raise_value_error("__enter__")
+    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):
-    _raise_value_error("__exit__")
+    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):
@@ -472,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/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 69d6843f5..eea734420 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -81,6 +81,20 @@ def _promote_integer_dtype(dtype: DTypeLike) -> DTypeLike:
       return dtypes.int_
   return dtype
 
+def check_where(name: str, where: ArrayLike | None) -> Array | None:
+  if where is None:
+    return where
+  check_arraylike(name, where)
+  where_arr = lax_internal.asarray(where)
+  if where_arr.dtype != bool:
+    # Deprecation added 2024-12-05
+    deprecations.warn(
+      'jax-numpy-reduction-non-boolean-where',
+      f"jnp.{name}: where must be None or a boolean array; got dtype={where_arr.dtype}.",
+      stacklevel=2)
+    return where_arr.astype(bool)
+  return where_arr
+
 
 ReductionOp = Callable[[Any, Any], Any]
 
@@ -101,6 +115,7 @@ def _reduction(a: ArrayLike, name: str, op: ReductionOp, init_val: ArrayLike,
   if out is not None:
     raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported.")
   check_arraylike(name, a)
+  where_ = check_where(name, where_)
   dtypes.check_user_dtype_supported(dtype, name)
   axis = core.concrete_or_error(None, axis, f"axis argument to jnp.{name}().")
 
@@ -730,6 +745,8 @@ def _logsumexp(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.logaddexp.reduce is not supported.")
   dtypes.check_user_dtype_supported(dtype, "jnp.logaddexp.reduce")
+  check_arraylike("logsumexp", a)
+  where = check_where("logsumexp", where)
   a_arr, = promote_dtypes_inexact(a)
   pos_dims, dims = _reduction_dims(a_arr, axis)
   amax = max(a_arr.real, axis=dims, keepdims=keepdims, where=where, initial=-np.inf)
@@ -748,6 +765,8 @@ def _logsumexp2(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.logaddexp2.reduce is not supported.")
   dtypes.check_user_dtype_supported(dtype, "jnp.logaddexp2.reduce")
+  check_arraylike("logsumexp2", a)
+  where = check_where("logsumexp2", where)
   ln2 = float(np.log(2))
   if initial is not None:
     initial *= ln2
@@ -850,6 +869,7 @@ def _mean(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
           upcast_f16_for_computation: bool = True,
           where: ArrayLike | None = None) -> Array:
   check_arraylike("mean", a)
+  where = check_where("mean", where)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.mean is not supported.")
 
@@ -1087,6 +1107,7 @@ def _var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
          out: None = None, correction: int | float = 0, keepdims: bool = False, *,
          where: ArrayLike | None = None) -> Array:
   check_arraylike("var", a)
+  where = check_where("var", where)
   dtypes.check_user_dtype_supported(dtype, "var")
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.var is not supported.")
@@ -1224,6 +1245,7 @@ def _std(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
          out: None = None, correction: int | float = 0, keepdims: bool = False, *,
          where: ArrayLike | None = None) -> Array:
   check_arraylike("std", a)
+  where = check_where("std", where)
   dtypes.check_user_dtype_supported(dtype, "std")
   if dtype is not None and not dtypes.issubdtype(dtype, np.inexact):
     raise ValueError(f"dtype argument to jnp.std must be inexact; got {dtype}")
@@ -1330,13 +1352,15 @@ def count_nonzero(a: ArrayLike, axis: Axis = None,
 
 def _nan_reduction(a: ArrayLike, name: str, jnp_reduction: Callable[..., Array],
                    init_val: ArrayLike, nan_if_all_nan: bool,
-                   axis: Axis = None, keepdims: bool = False, **kwargs) -> Array:
+                   axis: Axis = None, keepdims: bool = False, where: ArrayLike | None = None,
+                   **kwargs) -> Array:
   check_arraylike(name, a)
+  where = check_where(name, where)
   if not dtypes.issubdtype(dtypes.dtype(a), np.inexact):
-    return jnp_reduction(a, axis=axis, keepdims=keepdims, **kwargs)
+    return jnp_reduction(a, axis=axis, keepdims=keepdims, where=where, **kwargs)
 
   out = jnp_reduction(_where(lax_internal._isnan(a), _reduction_init_val(a, init_val), a),
-                      axis=axis, keepdims=keepdims, **kwargs)
+                      axis=axis, keepdims=keepdims, where=where, **kwargs)
   if nan_if_all_nan:
     return _where(all(lax_internal._isnan(a), axis=axis, keepdims=keepdims),
                   _lax_const(a, np.nan), out)
@@ -1755,6 +1779,7 @@ def nanmean(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None, out
     Array([[nan, nan, nan, nan]], dtype=float32)
   """
   check_arraylike("nanmean", a)
+  where = check_where("nanmean", where)
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.nanmean is not supported.")
   if dtypes.issubdtype(dtypes.dtype(a), np.bool_) or dtypes.issubdtype(dtypes.dtype(a), np.integer):
@@ -1848,6 +1873,7 @@ def nanvar(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None, out:
            [4.  ]], dtype=float32)
   """
   check_arraylike("nanvar", a)
+  where = check_where("nanvar", where)
   dtypes.check_user_dtype_supported(dtype, "nanvar")
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.nanvar is not supported.")
@@ -1943,6 +1969,7 @@ def nanstd(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None, out:
     Array([[0.5, 0.5, 0. , 0. ]], dtype=float32)
   """
   check_arraylike("nanstd", a)
+  where = check_where("nanstd", where)
   dtypes.check_user_dtype_supported(dtype, "nanstd")
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.nanstd is not supported.")
diff --git a/jax/_src/pallas/mosaic/verification.py b/jax/_src/pallas/mosaic/verification.py
index 08ff58770..61caa4087 100644
--- a/jax/_src/pallas/mosaic/verification.py
+++ b/jax/_src/pallas/mosaic/verification.py
@@ -145,15 +145,6 @@ class PrintCtx:
     self.level -= 1
     self.locals.append(self._indent(end) + "\n")
 
-  @contextlib.contextmanager
-  def comment_if_emitted(self, comment):
-    self.comment(comment)
-    yield
-    self.comment(comment)
-    if self.locals[-1] == self.locals[-2]:
-      self.locals.pop()
-    self.locals.pop()
-
   def get(self, value: ir.Value, default: Any = _UNSPECIFIED):
     if default is _UNSPECIFIED:
       return self.env[value]
@@ -367,17 +358,6 @@ def _print_op(ctx, op):
       return bin_op(ctx, "int", "%", *op.operands)
     case "arith.divsi":
       return bin_op(ctx, "int", "/", *op.operands)
-    case "arith.andi":
-      return bin_op(ctx, _model_type(op.result.type), "&", *op.operands)
-    case "arith.select":
-      cond, if_true, if_false = map(lambda o: ctx.get(o, None), op.operands)
-      if cond is None or if_true is None or if_false is None:
-        return NotImplemented
-      result_ty = _model_type(op.result.type)
-      return ctx.emit(result_ty, f"({cond} -> {if_true} : {if_false})")
-    case "arith.index_cast":
-      model = ctx.get(op.operands[0], None)
-      return ctx.emit("int", model) if model is not None else NotImplemented
     case "arith.cmpi":
       match op.predicate.value:
         case arith.CmpIPredicate.eq:
@@ -406,44 +386,12 @@ def _print_op(ctx, op):
         read_refs.append(model)
       with ctx.block("d_step {", "}"):  # Start reading
         for r in read_refs:
-          for loc in r.written_at(None):
-            ctx.emit(None, f"assert(!{loc})")
           for loc in r.readers_at(None):
             ctx.emit(None, f"{loc}++")
       with ctx.block("d_step {", "}"):  # Stop reading
         for r in read_refs:
           for loc in r.readers_at(None):
             ctx.emit(None, f"{loc}--")
-    case "vector.load":
-      ref = ctx.get(op.operands[0])
-      assert isinstance(ref, GlobalRefModel)
-      if (first_idx := ctx.get(op.operands[1], None)) is not None:
-        leading_load_len = ir.VectorType(op.result.type).shape[0]
-        ref = GlobalRefModel(f"{ref.base} + {first_idx}", leading_load_len)
-      with ctx.block("d_step {", "}"):  # Start reading
-        for loc in ref.written_at(None):
-          ctx.emit(None, f"assert(!{loc})")
-        for loc in ref.readers_at(None):
-          ctx.emit(None, f"{loc}++")
-      with ctx.block("d_step {", "}"):  # Stop reading
-        for loc in ref.readers_at(None):
-          ctx.emit(None, f"{loc}--")
-      return NotImplemented  # We don't model the result of the load.
-    case "vector.store":
-      ref = ctx.get(op.operands[1])  # Stored value goes first
-      assert isinstance(ref, GlobalRefModel)
-      if (first_idx := ctx.get(op.operands[2], None)) is not None:
-        leading_store_len = ir.VectorType(op.operands[0].type).shape[0]
-        ref = GlobalRefModel(f"{ref.base} + {first_idx}", leading_store_len)
-      with ctx.block("d_step {", "}"):  # Start writing
-        for loc in ref.readers_at(None):
-          ctx.emit(None, f"assert(!{loc})")
-        for loc in ref.written_at(None):
-          ctx.emit(None, f"assert(!{loc})")
-          ctx.emit(None, f"{loc} = 1")
-      with ctx.block("d_step {", "}"):  # Stop reading
-        for loc in ref.written_at(None):
-          ctx.emit(None, f"{loc} = 0")
     case "scf.for":
       carrys = [
           ctx.emit("int", ctx.get(arg))
@@ -471,7 +419,6 @@ def _print_op(ctx, op):
               ctx.emit(None, f"{c} = {ctx.get(new)}")
           ctx.emit(None, f"{induction_var} = {induction_var} + {step}")
         ctx.emit(None, ":: else -> break")
-      ctx.emit(None, "skip")  # To avoid "Jump into d_step sequence errors"
       if len(carrys) == 1:
         return carrys[0]
       else:
@@ -503,27 +450,16 @@ def bin_op(ctx, result_ty, op, lhs, rhs):
   return ctx.emit(result_ty, f"{lhs} {op} {rhs}")
 
 
-def _model_type(ty):
-  if ir.IntegerType.isinstance(ty):
-    if ir.IntegerType(ty).width == 1:
-      return "bool"
-    else:
-      return "int"
-  else:
-    raise NotImplementedError(ty)
-
-
 def _print_block(ctx, block):
   for op in block:
     try:
-      with ctx.comment_if_emitted(op.OPERATION_NAME):
-        results = _print_op(ctx, op)
+      results = _print_op(ctx, op)
     except Exception as e:
       raise RuntimeError(f"Failed to print op: {op}") from e
     if results is NotImplemented:
       continue
     if not op.results:
-      assert results is None or results == ()
+      assert results is None
     elif len(op.results) > 1:
       raise NotImplementedError(op)
     else:
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index e04b0ff1a..2900da133 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -17,13 +17,13 @@
 from __future__ import annotations
 
 import collections
-from collections.abc import Hashable, MutableMapping, MutableSequence, Sequence
+from collections.abc import MutableMapping, MutableSequence, Sequence
 import contextlib
 import dataclasses
 import functools
 import itertools as it
 import math
-from typing import Any, Protocol, cast
+from typing import Any, Hashable, Protocol, cast
 
 import jax
 from jax import lax
@@ -192,7 +192,7 @@ def _reduce_sum_resource_estimator(x_aval: jax_core.ShapedArray, *, axes) -> int
 @dataclasses.dataclass
 class ModuleContext:
   name: str
-  grid_names: Sequence[Hashable] | None
+  grid_mapping: pallas_core.GridMapping
   program_ids: Sequence[ir.Value] | None
   approx_math: bool
   runtime_smem: ir.Value  # ir.MemRefType
@@ -517,7 +517,7 @@ def lower_jaxpr_to_module(
       grouped_barriers[barrier].append(barrier_ref)
     module_ctx = ModuleContext(
         name_and_src_info.name,
-        grid_mapping.grid_names,
+        grid_mapping,
         None,
         approx_math,
         runtime_smem,
@@ -1046,6 +1046,8 @@ def _get_lowering_rule(ctx: LoweringRuleContext, x_smem, *leaves, tree):
   x_smem, transforms = _handle_reshaping(x_smem, transforms)
   x_smem, transforms = _handle_indexing(x_smem, transforms)
 
+  print("ctx:", ctx)
+  print("transforms:", transforms)
   match transforms:
     case (gpu_core.UnswizzleRef(swizzle), gpu_core.UntileRef(tiling)):
       if tiling != (64, swizzle // x_aval.dtype.itemsize):
@@ -1290,7 +1292,7 @@ def _reduce_max_lowering_rule(ctx: LoweringRuleContext, x, *, axes):
 @register_lowering_rule(lax.axis_index_p)
 def _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
   i32 = ir.IntegerType.get_signless(32)
-  grid_names = ctx.module_ctx.grid_names
+  grid_names = ctx.module_ctx.grid_mapping.grid_names
   squashed_dims = ctx.module_ctx.squashed_dims
   if squashed_dims:
     unsquashed_names = grid_names[-3:]
diff --git a/jax/_src/pallas/mosaic_gpu/pipeline.py b/jax/_src/pallas/mosaic_gpu/pipeline.py
index ee3f03f18..069b8d9e7 100644
--- a/jax/_src/pallas/mosaic_gpu/pipeline.py
+++ b/jax/_src/pallas/mosaic_gpu/pipeline.py
@@ -181,6 +181,17 @@ def emit_pipeline(
     delay_release = 0  # No need to delay anything.
 
   def pipeline(*gmem_refs: pallas_core.AbstractMemoryRef):
+    for gmem_ref, spec in zip(gmem_refs, it.chain(in_specs, out_specs)):
+      if any(
+          spec.block_shape[-idx] * grid[-idx] != gmem_ref.shape[-idx]  # type: ignore
+          for idx in range(1, len(grid) + 1)
+          if spec.block_shape is not None
+      ):
+        raise NotImplementedError(
+            f"Cannot emit a pipeline over the {grid=} for {gmem_ref} with block"
+            f" shape {spec.block_shape}."
+        )
+
     in_gmem_refs, out_gmem_refs = util.split_list(gmem_refs, [len(in_specs)])
     in_smem_refs, out_smem_refs = util.split_list(
         [
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 196dd8b01..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')
@@ -2385,7 +2385,7 @@ def _pjit_transpose(cts_in, *primals_in,
     _set_states(attrs_tracked, final_states)
 
   return tree_unflatten(cts_out_treedef, nz_cts_out)
-ad.primitive_transposes[pjit_p] = _pjit_transpose
+ad.reducing_transposes[pjit_p] = _pjit_transpose
 
 
 @weakref_lru_cache
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/mosaic/gpu/core.py b/jax/experimental/mosaic/gpu/core.py
index d8774e932..4b4882f65 100644
--- a/jax/experimental/mosaic/gpu/core.py
+++ b/jax/experimental/mosaic/gpu/core.py
@@ -806,6 +806,16 @@ def _launch(
         )
     )
 
+    smem_ref_tree = _construct_smem_reftree(
+        cluster, dynamic_smem, smem_buffers
+    )
+    # TODO(apaszke): Skip the following if no barriers were initialized.
+    nvvm.fence_mbarrier_init()
+    if math.prod(cluster) != 1:
+      nvvm.cluster_arrive_relaxed(aligned=ir.UnitAttr.get())
+      nvvm.cluster_wait(aligned=ir.UnitAttr.get())
+    gpu.barrier()
+
     if profiler_spec:
       prof_smem = memref.view(
           ir.MemRefType.get(
@@ -822,19 +832,7 @@ def _launch(
 
     ptr_ty = ir.Type.parse("!llvm.ptr")
     scratch_ptr = builtin.unrealized_conversion_cast([ptr_ty], [scratch_arr])
-    ctx = LaunchContext(launch_op, scratch_ptr, cluster, prof)
-    with ctx.named_region("Init"):
-      smem_ref_tree = _construct_smem_reftree(
-          cluster, dynamic_smem, smem_buffers
-      )
-      # TODO(apaszke): Skip the following if no barriers were initialized.
-      nvvm.fence_mbarrier_init()
-      if math.prod(cluster) != 1:
-        nvvm.cluster_arrive_relaxed(aligned=ir.UnitAttr.get())
-        nvvm.cluster_wait(aligned=ir.UnitAttr.get())
-      gpu.barrier()
-
-    yield ctx, smem_ref_tree
+    yield LaunchContext(launch_op, scratch_ptr, cluster, prof), smem_ref_tree
     if prof is not None:
       prof.finalize(grid=grid, block=block)
     gpu.terminator()
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index f6cab5654..2279df4f3 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -673,17 +673,33 @@ class BarrierRef:
         1,
     )
 
-  def wait_parity(self, parity):
+  def wait_parity(self, parity, expect_wait=False):
+    i1 = ir.IntegerType.get_signless(1)
     i32 = ir.IntegerType.get_signless(32)
-    ticks = arith.constant(i32, 10000000)
+    ticks = c(10000000, i32)
+    address = self.get_ptr()
     parity = arith.extui(i32, parity)
-    nvvm.mbarrier_try_wait_parity_shared(self.get_ptr(), parity, ticks)
+    if expect_wait:
+      nvvm.mbarrier_try_wait_parity_shared(address, parity, ticks)
+      return
+    barrier_ready = llvm.inline_asm(
+        i1,
+        [address, parity],
+        "mbarrier.test_wait.parity.shared.b64 $0, [$1], $2;",
+        "=b,l,r",
+        has_side_effects=True,
+    )
+    should_wait = arith.xori(barrier_ready, c(1, i1))
+    should_wait = llvm.intr_expect(should_wait, c(0, i1))
+    with ir.InsertionPoint(scf.IfOp(should_wait).then_block):
+      nvvm.mbarrier_try_wait_parity_shared(address, parity, ticks)
+      scf.yield_([])
 
-  def wait(self):
+  def wait(self, expect_wait=False):
     parities = memref.load(self.phases, [])
     parity, new_parities = self.update_parities(parities)
     memref.store(new_parities, self.phases, [])
-    self.wait_parity(parity)
+    self.wait_parity(parity, expect_wait=expect_wait)
 
   def update_parities(self, parities: ir.Value) -> tuple[ir.Value, ir.Value]:
     i32 = ir.IntegerType.get_signless(32)
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:
diff --git a/jaxlib/mosaic/dialect/tpu/tpu.td b/jaxlib/mosaic/dialect/tpu/tpu.td
index c5142f48d..5dad1309a 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu.td
+++ b/jaxlib/mosaic/dialect/tpu/tpu.td
@@ -342,13 +342,8 @@ def TPU_RepeatOp : TPU_Op<"repeat", [Pure]> {
 }
 
 def TPU_BroadcastInSublanesOp : TPU_Op<"broadcast_in_sublanes", [Pure]> {
-  let description = [{
-    For each sublane `i`, broadcasts the value in lane `lane + i` along the entire
-    sublane. If `lane + i` is not in [0, lane_count), then the value in sublane `i`
-    is not defined (can be anything).
-  }];
   let arguments = (ins
-    TPU_Vreg:$source,  // All sublanes should be equal.
+    AnyVectorOfNonZeroRank:$source,  // All sublanes should be equal.
     I32Attr:$lane  // Coordinates of the first element to take.
   );
   // Output shape should be the same, except for position dim which contains
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_ops.cc b/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
index 07e2e3e19..2107cd7fc 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
+++ b/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
@@ -93,25 +93,6 @@ LogicalResult MemRefSliceOp::verify() {
   auto target_type = getType();
   auto target_layout = target_type.getLayout();
   auto target_memory_space = target_type.getMemorySpace();
-  auto indices = getBaseIdx();
-  auto slice_shape = getResult().getType().getShape();
-  if (!source_type.hasStaticShape()) {
-    return emitOpError(
-        "Only slicing of memrefs with static shapes is supported.");
-  }
-  auto source_shape = source_type.getShape();
-  bool is_semaphore =
-      HasMemorySpace(source_type, tpu::MemorySpace::kSemaphoreMem);
-  if (is_semaphore &&
-      !isa<SemaphoreType, DMASemaphoreType>(source_type.getElementType())) {
-    return emitOpError(
-        "References to semaphore memory space must have a semaphore element "
-        "type.");
-  }
-  if (indices.size() != slice_shape.size() ||
-      indices.size() != source_shape.size()) {
-    return emitOpError("Indices and slice shapes must match.");
-  }
   // TODO(apaszke): Check that the result has a smaller shape.
   // TODO(apaszke): Check that strides are equivalent.
   // Source and target attributes may be different before propagation is done by
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
index 50a7d5734..cac197479 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -655,105 +655,6 @@ FailureOr<SmallVector<Layout>> getInLayouts(
   return in_layouts;
 }
 
-// Insert a minor dimension to the implicit shape. The original minor dimension
-// becomes the new second minor dimension, laid out across sublanes.
-//
-// The returned vreg array uses the original tiling and the offsets specified in
-// new_offsets to hold the value with the new implicit shape.
-//
-// Args:
-// vregs:       The vreg array with *implicit* array shape.
-// ishape:      The implicit shape of the represented value.
-// layout:      The layout used for the represented value. The implicit
-//              dimension is ignored, since this function operates directly at
-//              the level of the implicit shape.
-// new_offsets: The offsets to use for the layout of the returned vreg array.
-FailureOr<xla::Array<Value>> insertImplicitMinorDimension(
-    RewriteContext &ctx, OpBuilder &builder, const Location loc,
-    const xla::Array<Value> &vregs, const ArrayRef<int64_t> ishape,
-    const VectorLayout &layout, const LayoutOffsets new_offsets) {
-  if (layout.bitwidth() != 32 || !layout.hasNativeTiling(ctx.target_shape)) {
-    return emitError(loc, "Not implemented: Unsupported bitwidth or tiling");
-  }
-  if (layout.offsets()[1].has_value()) {
-    if (!new_offsets[0]) {
-      // TODO(tlongeri): This can only be valid if the dim size is 1.
-      return emitError(loc, "Not implemented: Replication mismatch");
-    }
-    if (*new_offsets[0] != *layout.offsets()[1] % ctx.target_shape[0] &&
-        *layout.offsets()[1] + *(ishape.end() - 1) > ctx.target_shape[1]) {
-      // This requires blending data from different vregs.
-      return emitError(loc,
-                       "Not implemented: Misaligned offsets and shape does not "
-                       "fit in one vreg");
-    }
-  }
-  // new_layout is only to get the new vreg array shape, the implicit dim is
-  // irrelevant (since we already have the implicit shape):
-  const VectorLayout new_layout(layout.bitwidth(), new_offsets, layout.tiling(),
-                                VectorLayout::ImplicitDim::kNone);
-  SmallVector<int64_t> new_ishape(ishape);
-  new_ishape.push_back(1);
-  xla::Array<Value> new_vregs(new_layout.tileArrayShape(
-      /*src_is_implicit=*/true, /*res_is_implicit=*/true, std::move(new_ishape),
-      ctx.target_shape));
-  // Preallocate an indices vector to avoid repeated allocations:
-  SmallVector<int64_t> idxs;
-  new_vregs.Each([&](const absl::Span<const int64_t> dst_idx,
-                     Value *const dst_vreg) {
-    // Indices of the new vreg in the new vreg array:
-    const int64_t new_2nd_minor_idx = *(dst_idx.end() - 2);
-    const int64_t new_3rd_minor_idx = *(dst_idx.end() - 3);
-    idxs.assign(dst_idx.begin(), dst_idx.end());
-    if (!layout.offsets()[0].has_value() && new_3rd_minor_idx != 0) {
-      // All vregs along that dimension are the same
-      *(idxs.end() - 3) = 0;
-      *dst_vreg = new_vregs(idxs);
-    } else if (!layout.offsets()[1].has_value() && new_2nd_minor_idx != 0) {
-      // All vregs along that dimension are the same
-      *(idxs.end() - 2) = 0;
-      *dst_vreg = new_vregs(idxs);
-    } else {
-      // dst_vreg will hold slice [row_idx, col_idx:(col_idx + target_shape[0])]
-      // of the after-offsets source shape
-      const int64_t row_idx =
-          layout.offsets()[0] ? new_3rd_minor_idx + *layout.offsets()[0] : 0;
-      const int64_t col_idx = layout.offsets()[1]
-                                  ? new_2nd_minor_idx * ctx.target_shape[0] +
-                                        *layout.offsets()[1] - *new_offsets[0]
-                                  : 0;
-
-      idxs.pop_back();
-      *(idxs.end() - 2) = row_idx / ctx.target_shape[0];
-      *(idxs.end() - 1) = col_idx / ctx.target_shape[1];
-      Value src_vreg = vregs(idxs);
-      // TODO(tlongeri): We can sometimes skip operations when dst_vreg will
-      // hold a single non-padding element (first or last) and we don't need
-      // replication in the output.
-      if (layout.offsets()[0].has_value()) {
-        // [ . . . . . . . . ]    [ . . . . a b c d ]
-        // [ . . . . a b c d ] => [ . . . . a b c d ]
-        // [ . . . . . . . . ]    [ . . . . a b c d ]
-        // [ . . . . . . . . ]    [ . . . . a b c d ]
-        src_vreg = broadcastSublane(
-            builder, src_vreg,
-            /*sublane_idx=*/row_idx % ctx.target_shape[0], ctx.target_shape);
-      }
-      if (layout.offsets()[1].has_value()) {
-        // [ . . . . a b c d ]    [ a a a a a a a a ]
-        // [ . . . . a b c d ] => [ b b b b b b b b ]
-        // [ . . . . a b c d ]    [ c c c c c c c c ]
-        // [ . . . . a b c d ]    [ d d d d d d d d ]
-        src_vreg = builder.create<BroadcastInSublanesOp>(
-            loc, src_vreg.getType(), src_vreg,
-            /*lane=*/col_idx % ctx.target_shape[1]);
-      }
-      *dst_vreg = src_vreg;
-    }
-  });
-  return new_vregs;
-}
-
 LogicalResult elementwise_op_rule(RewriteContext &ctx, Operation &op,
                                   const ArrayRef<Layout> layouts_in,
                                   const ArrayRef<Layout> layouts_out) {
@@ -4254,16 +4155,54 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
                layout_in.bitwidth() == 32 &&
                layout_in.hasNativeTiling(ctx.target_shape) &&
                layout_in.tiling() == layout_out.tiling() &&
-               (!layout_in.offsets()[1].has_value() ||
-                *layout_in.offsets()[1] % ctx.target_shape[0] ==
-                    layout_out.offsets()[0] ||
-                *layout_in.offsets()[1] + src_tiled_dims[1] <=
-                    ctx.target_shape[1])) {
-      FAILUREOR_ASSIGN_OR_RETURN(
-          xla::Array<Value> dst_vregs_local,
-          insertImplicitMinorDimension(ctx, builder, op.getLoc(), src_vregs,
-                                       layout_in.implicitShape(src_shape),
-                                       layout_in, layout_out.offsets()));
+               layout_in.offsets()[0].value_or(0) == 0 &&
+               layout_in.offsets()[1] == 0 && layout_out.offsets()[0] == 0
+               // layout_out.offsets[1] can be anything, as we produce a
+               // replicated result
+    ) {
+      // First, insert the new singleton lane dimension.
+      SmallVector<int64_t> s = layout_in.implicitShape(src_shape);
+      s.push_back(1);
+      xla::Array<Value> dst_vregs_local(layout_out.tileArrayShape(
+          /*src_is_implicit=*/true, /*res_is_implicit=*/true, std::move(s),
+          ctx.target_shape));
+      TPU_ASSERT_EQ_OP(dst_vregs_local.dimensions().back(),
+                       1);  // We're inserting a singleton dimension
+      dst_vregs_local.Each(
+          [&](const absl::Span<const int64_t> dst_idx, Value *const dst_vreg) {
+            const int64_t col_idx = *(dst_idx.end() - 2);
+            const int64_t row_idx = *(dst_idx.end() - 3);
+            auto [sublanes_in_lane, rem] =
+                std::div(ctx.target_shape[1], ctx.target_shape[0]);
+            CHECK_EQ(rem, 0);
+            if (!layout_in.offsets()[0].has_value() && row_idx != 0) {
+              return;  // All vregs along that dimension are the same.
+            }
+            SmallVector<int64_t> src_idx(toArrayRef(dst_idx));
+            src_idx.pop_back();
+            *(src_idx.end() - 2) /= ctx.target_shape[0];
+            *(src_idx.end() - 1) /= sublanes_in_lane;
+            Value col_vreg = src_vregs(src_idx);
+            // BroadcastInSublanesOp requires the sublanes to be replicated.
+            if (layout_in.offsets()[0].has_value()) {
+              const int32_t sublane = row_idx % ctx.target_shape[0];
+              col_vreg = broadcastSublane(builder, col_vreg, sublane,
+                                          ctx.target_shape);
+            }
+            *dst_vreg = builder.create<BroadcastInSublanesOp>(
+                col_vreg.getType(), col_vreg,
+                /*lane=*/(col_idx % sublanes_in_lane) * ctx.target_shape[0]);
+          });
+      if (!layout_in.offsets()[0].has_value()) {
+        // Broadcast the sublane vregs.
+        // TODO(tlongeri): This could be done more efficiently
+        dst_vregs_local.Each([&](const absl::Span<const int64_t> dst_idx,
+                                 Value *const dst_vreg) {
+          SmallVector<int64_t> first_row_idx(toArrayRef(dst_idx));
+          *(first_row_idx.end() - 3) = 0;
+          *dst_vreg = dst_vregs_local(first_row_idx);
+        });
+      }
       // Now, reshape the major axes of the vreg array.
       dst_vregs_local.Reshape(
           layout_out.tileArrayImplicitShape(dst_shape, ctx.target_shape));
@@ -6177,49 +6116,35 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
   }
   const int packing = src.packing();
   const int8_t bitwidth = src.bitwidth();
+  VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
+                   src.implicit_dim());
+  if (!dst.isValid(target_shape)) {
+    return emitError(loc, "Not implemented: invalid offsets in tiling target");
+  }
+  auto dst_tiles_shape =
+      dst.tileArrayImplicitShape(vty.getShape(), target_shape);
   // Handle retiling from (1, 128) to (8, 128) for 32-bit data with replicating
   // sublanes.
   if (try_replicate_rows && packing == 1 &&
       *(vregs.dimensions().end() - 2) == 1 &&
+      src.offsets() == LayoutOffsets{0, 0} &&
       src.tiling() == std::array<int64_t, 2>{1, ctx.target_shape[1]} &&
       dst_tiling == ctx.target_shape) {
-    DCHECK_EQ(src.offsets()[0].value_or(0), 0);
-    const LayoutOffset dst_minor_offset =
-        src.offsets()[1] ? LayoutOffset(*src.offsets()[1] % target_shape[1])
-                         : std::nullopt;
-    const VectorLayout dst(bitwidth, {std::nullopt, dst_minor_offset},
-                           dst_tiling, src.implicit_dim());
-    xla::Array<Value> retiled(
-        dst.tileArrayImplicitShape(vty.getShape(), target_shape));
+    xla::Array<Value> retiled(dst_tiles_shape);
     retiled.Each([&](absl::Span<const int64_t> idx, Value *tile) {
       SmallVector<int64_t> src_idx(idx.begin(), idx.end());
       *(src_idx.end() - 2) *= target_shape[0];
-      if (!src.offsets()[1].has_value()) {
-        // With (1, 128) tiling each vreg holds values from a single row. This
-        // means that if the columns are replicated, then the whole vreg is
-        // already replicated.
-        *(src_idx.end() - 1) = 0;
-        *tile = vregs(src_idx);
-      } else {
-        // The column (in units of sublanes) of the sublane we want:
-        const int64_t sublane_column =
-            *(src_idx.end() - 1) + *src.offsets()[1] / target_shape[1];
-        *(src_idx.end() - 1) = sublane_column / target_shape[0];
-        const int64_t src_sl_idx = sublane_column % target_shape[0];
-        *tile =
-            broadcastSublane(builder, vregs(src_idx), src_sl_idx, target_shape);
-      }
+      *(src_idx.end() - 1) /= target_shape[0];
+      const int64_t src_sl_idx = *(idx.end() - 1) % target_shape[0];
+      CHECK_EQ(src.getImplicitTiledDims(vty.getShape(), 1)[0], 1);
+      *tile =
+          broadcastSublane(builder, vregs(src_idx), src_sl_idx, target_shape);
     });
-    // We have successfully replicated sublanes
+    // We have successfully replicated sublanes.
+    dst = VectorLayout(bitwidth, {std::nullopt, dst.offsets()[1]}, dst_tiling,
+                       dst.implicit_dim());
     return std::pair(dst, std::move(retiled));
   }
-  VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
-                   src.implicit_dim());
-  if (!dst.isValid(target_shape)) {
-    return emitError(loc, "Not implemented: invalid offsets in tiling target");
-  }
-  auto dst_tiles_shape =
-      dst.tileArrayImplicitShape(vty.getShape(), target_shape);
   // (8,128) -> (8 * packing,128) tiling change for packed type.
   if (bitwidth < 32 && 32 % bitwidth == 0 && src_tiling == ctx.target_shape &&
       dst_tiling == std::array<int64_t, 2>{ctx.target_shape[0] * dst.packing(),
@@ -6431,26 +6356,6 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeImplicitDim(
     });
     return std::make_pair(dst, new_vregs);
   }
-  if (src.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
-      dst_implicit_dim == VectorLayout::ImplicitDim::kMinor &&
-      src.bitwidth() == 32 && src.hasNativeTiling(ctx.target_shape)) {
-    // TODO(tlongeri): Make insertImplicitMinorDimension more flexible about
-    //                 offsets, then we can pass dst_offset_hints directly.
-    const LayoutOffset dst_2nd_minor_offset =
-        !src.offsets()[1] || *src.offsets()[1] + *(vty.getShape().end() - 1) <=
-                                 ctx.target_shape[1]
-            ? dst_offset_hints[0]
-            : LayoutOffset(*src.offsets()[1] % ctx.target_shape[0]);
-    VectorLayout dst(src.bitwidth(),
-                     {dst_2nd_minor_offset, dst_offset_hints[1]}, src.tiling(),
-                     VectorLayout::ImplicitDim::kMinor);
-    FAILUREOR_ASSIGN_OR_RETURN(
-        xla::Array<Value> dst_vregs,
-        insertImplicitMinorDimension(ctx, builder, loc, vregs,
-                                     src.implicitShape(vty.getShape()), src,
-                                     dst.offsets()));
-    return std::make_pair(dst, std::move(dst_vregs));
-  }
   return emitError(loc,
                    "Not implemented: Unsupported implicit dim change: from ")
          << src << " to " << dst_implicit_dim;
diff --git a/jaxlib/mosaic/dialect/tpu/util.cc b/jaxlib/mosaic/dialect/tpu/util.cc
index c5f983376..d434837ef 100644
--- a/jaxlib/mosaic/dialect/tpu/util.cc
+++ b/jaxlib/mosaic/dialect/tpu/util.cc
@@ -141,10 +141,4 @@ bool canReinterpretToUntiledMemref(TypedValue<MemRefType> tiled_memref,
   return *(tiled_layout.getTileStrides().end() - 1) == 1 &&
          *(tiled_layout.getTileStrides().end() - 2) == 1;
 }
-
-bool HasMemorySpace(MemRefType ty, tpu::MemorySpace space) {
-  auto memory_space =
-      dyn_cast_or_null<tpu::MemorySpaceAttr>(ty.getMemorySpace());
-  return memory_space && memory_space.getValue() == space;
-}
 }  // namespace mlir::tpu
diff --git a/jaxlib/mosaic/dialect/tpu/util.h b/jaxlib/mosaic/dialect/tpu/util.h
index 9052afad4..7c602e9a0 100644
--- a/jaxlib/mosaic/dialect/tpu/util.h
+++ b/jaxlib/mosaic/dialect/tpu/util.h
@@ -115,8 +115,6 @@ bool canReinterpretToUntiledMemref(TypedValue<MemRefType> tiled_memref,
                                    const std::array<int64_t, 2> &target_shape,
                                    bool allow_minormost_padding = false);
 
-// Determines whether the given MemRefType has the given memory space.
-bool HasMemorySpace(MemRefType ty, tpu::MemorySpace space);
 }  // namespace mlir::tpu
 
 #endif  // THIRD_PARTY_PY_JAX_JAXLIB_MOSAIC_DIALECT_TPU_UTIL_H_
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index db34354f4..232d19cfb 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "30f22f4d6cb523e035c237f30aeac4e00ae34821"
-XLA_SHA256 = "546dc97a5bee684b3baf1c14c00ef6c73f18c717ebb97c000a35f683bf53c244"
+XLA_COMMIT = "05f004e8368c955b872126b1c978c60e33bbc5c8"
+XLA_SHA256 = "f0bedada96f5f1d09f5047c7f9db32d460d147bd0f192607cfbbee9fe5ee2d5f"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index be6208e6e..2bef35fbd 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -448,6 +448,34 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, atol=tol, rtol=tol)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @jtu.sample_product(rec=JAX_REDUCER_INITIAL_RECORDS)
+  def testReducerWhereNonBooleanErrorInitial(self, rec):
+    dtype = rec.dtypes[0]
+    x = jnp.zeros((10,), dtype)
+    where = jnp.ones(10, dtype=int)
+    func = getattr(jnp, rec.name)
+    def assert_warns_or_errors(msg):
+      if deprecations.is_accelerated("jax-numpy-reduction-non-boolean-where"):
+        return self.assertRaisesRegex(ValueError, msg)
+      else:
+        return self.assertWarnsRegex(DeprecationWarning, msg)
+    with assert_warns_or_errors(f"jnp.{rec.name}: where must be None or a boolean array"):
+      func(x, where=where, initial=jnp.array(0, dtype=dtype))
+
+  @jtu.sample_product(rec=JAX_REDUCER_WHERE_NO_INITIAL_RECORDS)
+  def testReducerWhereNonBooleanErrorNoInitial(self, rec):
+    dtype = rec.dtypes[0]
+    x = jnp.zeros((10,), dtype)
+    where = jnp.ones(10, dtype=int)
+    func = getattr(jnp, rec.name)
+    def assert_warns_or_errors(msg):
+      if deprecations.is_accelerated("jax-numpy-reduction-non-boolean-where"):
+        return self.assertRaisesRegex(ValueError, msg)
+      else:
+        return self.assertWarnsRegex(DeprecationWarning, msg)
+    with assert_warns_or_errors(f"jnp.{rec.name}: where must be None or a boolean array"):
+      func(x, where=where)
+
   @parameterized.parameters(itertools.chain.from_iterable(
     jtu.sample_product_testcases(
       [dict(name=rec.name, rng_factory=rec.rng_factory, inexact=rec.inexact,
diff --git a/tests/memories_test.py b/tests/memories_test.py
index 9c9b3a4ad..ca676a2b1 100644
--- a/tests/memories_test.py
+++ b/tests/memories_test.py
@@ -1578,7 +1578,7 @@ class ComputeOffload(jtu.BufferDonationTestCase):
         test_fn,
         out_shardings=(
             Layout(custom_dll, sharding),
-            Layout(custom_dll_linear, p_sharding),
+            Layout(custom_dll, p_sharding),
         ),
     )
     x_out, y_out = jit_fn(x, y)
@@ -1621,7 +1621,7 @@ class ComputeOffload(jtu.BufferDonationTestCase):
         test_fn,
         out_shardings=(
             Layout(custom_dll, sharding),
-            Layout(custom_dll_linear, p_sharding),
+            Layout(custom_dll, p_sharding),
         ),
     )
     x_out, y_out = jit_fn(x, y)
diff --git a/tests/pallas/export_back_compat_pallas_test.py b/tests/pallas/export_back_compat_pallas_test.py
index 462597e56..1b810bcb6 100644
--- a/tests/pallas/export_back_compat_pallas_test.py
+++ b/tests/pallas/export_back_compat_pallas_test.py
@@ -48,10 +48,9 @@ class CompatTest(bctu.CompatTestBase):
       self.skipTest("Only works on GPUs with capability >= sm80")
     super().setUp()
 
-  @unittest.skip("This test is checking backwards compatibility "
+  @unittest.skip("TODO(necula): This test is checking backwards compatibility "
                  "of Triton IR, but Triton doesn't promise backwards "
-                 "compatibility for its IR, and we have since removed "
-                 "the corresponding custom call from the guaranteed stable list.")
+                 "compatibility for its IR.")
   def test_triton_add_one(self):
     def func(x):
       def add_one(x_ref, o_ref):
diff --git a/tests/pallas/export_pallas_test.py b/tests/pallas/export_pallas_test.py
index 8b18f706a..70e40e1f2 100644
--- a/tests/pallas/export_pallas_test.py
+++ b/tests/pallas/export_pallas_test.py
@@ -50,10 +50,6 @@ class ExportTest(jtu.JaxTestCase):
     exp = export.export(
         add_vectors,
         platforms=["tpu", "cuda"],
-        # The Pallas GPU custom call is not enabled for export by default.
-        disabled_checks=[
-            export.DisabledSafetyCheck.custom_call("__gpu$xla.gpu.triton")
-        ]
     )(a, a)
 
     if (jtu.device_under_test() == "tpu" or
diff --git a/tests/pallas/tpu_pallas_distributed_test.py b/tests/pallas/tpu_pallas_distributed_test.py
index 7b3bd70ef..f7e3eaaf0 100644
--- a/tests/pallas/tpu_pallas_distributed_test.py
+++ b/tests/pallas/tpu_pallas_distributed_test.py
@@ -15,8 +15,6 @@
 """Tests for distributed pallas TPU operations."""
 
 import functools
-import os
-import tempfile
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -515,62 +513,5 @@ class PallasCallRemoteDMAInterpretTest(parameterized.TestCase):
                                atol=1e-5,
                                rtol=1e-3)
 
-
-class VerificationTest(jtu.JaxTestCase):
-
-  def test_verification(self):
-    if (num_devices := jax.local_device_count()) <= 1:
-      self.skipTest('Test requires multiple devices.')
-    if not jtu.is_device_tpu_at_least(4) or jax.devices()[0].num_cores > 1:
-      self.skipTest('Test requires a new single-core TPU.')
-    def kernel_body(in_ref, out_ref, scratch_ref, send_sem, recv_sem, capacity_sem):
-      my_id = lax.axis_index('x')
-      dst_id = jnp.where(my_id == num_devices - 1, 0, my_id + 1)
-      src_id = jnp.where(my_id == 0, num_devices - 1, my_id - 1)
-      pltpu.semaphore_signal(capacity_sem, 1, device_id=src_id)
-      out_ref[...] = jnp.zeros_like(out_ref)
-      scratch_ref[0] = in_ref[0]
-
-      @functools.partial(lax.fori_loop, 0, num_devices - 1, init_val=None)
-      def _(i, _):
-        slot = i % 2
-        next_slot = 1 - slot
-        pltpu.semaphore_wait(capacity_sem, 1)
-        copy = pltpu.async_remote_copy(
-            scratch_ref.at[slot],
-            scratch_ref.at[next_slot],
-            send_sem,
-            recv_sem,
-            device_id=dst_id,
-        )
-        out_ref[...] += scratch_ref[slot]
-        copy.wait()
-        pltpu.semaphore_signal(capacity_sem, 1, device_id=src_id)
-      out_ref[...] += scratch_ref[(num_devices - 1) % 2]
-      pltpu.semaphore_wait(capacity_sem, 1)
-
-    kernel = pl.pallas_call(
-        kernel_body,
-        out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float32),
-        scratch_shapes=[
-            pltpu.VMEM((2, 128, 128), jnp.float32),
-            pltpu.SemaphoreType.DMA,
-            pltpu.SemaphoreType.DMA,
-            pltpu.SemaphoreType.REGULAR,
-        ],
-    )
-    devices = mesh_utils.create_device_mesh((num_devices,))
-    mesh = jax.sharding.Mesh(devices, ['x'])
-    # This is just a smoke test to ensure that the verification does not crash.
-    with tempfile.TemporaryDirectory() as tmpdir:
-      previous_config = jax.config.read('jax_pallas_dump_promela_to')
-      jax.config.update('jax_pallas_dump_promela_to', tmpdir)
-      shard_map.shard_map(
-          kernel, mesh=mesh, in_specs=P('x'), out_specs=P(None), check_rep=False
-      )(jnp.ones((8, 128, 128), jnp.float32))
-      jax.config.update('jax_pallas_dump_promela_to', previous_config)
-      self.assertNotEmpty(os.listdir(tmpdir))
-
-
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
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')