Repository: google/jax
Pull request #22616 (2024-07-24 08:55:07)

DOC: Improved docs for `jnp.triu_indices` and `jnp.tril_indices`
Part of #21461

dfm: Looks good to me! Can you squash into one commit? Thanks!

Patch:
diff --git a/docs/pallas/grid_blockspec.md b/docs/pallas/grid_blockspec.md
index 521bf78a6..a5b27254f 100644
--- a/docs/pallas/grid_blockspec.md
+++ b/docs/pallas/grid_blockspec.md
@@ -34,7 +34,7 @@ This generalizes to any tuple of integers (a length `d` grid will correspond
 to `d` nested loops).
 The kernel is executed as many times
 as `prod(grid)`.
-The default grid value `()` results in one
+The default grid value `None` stands for `()`, and results in one
 kernel invocation.
 Each of these invocations is referred to as a "program".
 To access which program (i.e. which element of the grid) the kernel is currently
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index e4a2edc97..21c410787 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -4853,7 +4853,7 @@ def triu(m: ArrayLike, k: int = 0) -> Array:
 
   Args:
     m: input array. Must have ``m.ndim >= 2``.
-    k: k: optional, int, default=0. Specifies the sub-diagonal below which the
+    k: optional, int, default=0. Specifies the sub-diagonal below which the
       elements of the array are set to zero. ``k=0`` refers to main diagonal,
       ``k<0`` refers to sub-diagonal below the main diagonal and ``k>0`` refers
       to sub-diagonal above the main diagonal.
@@ -4953,8 +4953,57 @@ def _triu_size(n, m, k):
     return mk * (mk + 1) // 2 + mk * (m - k - mk)
 
 
-@util.implements(np.triu_indices)
 def triu_indices(n: int, k: int = 0, m: int | None = None) -> tuple[Array, Array]:
+  """Return the indices of upper triangle of an array of size ``(n, m)``.
+
+  JAX implementation of :func:`numpy.triu_indices`.
+
+  Args:
+    n: int. Number of rows of the array for which the indices are returned.
+    k: optional, int, default=0. Specifies the sub-diagonal on and above which
+      the indices of upper triangle are returned. ``k=0`` refers to main diagonal,
+      ``k<0`` refers to sub-diagonal below the main diagonal and ``k>0`` refers
+      to sub-diagonal above the main diagonal.
+    m: optional, int. Number of columns of the array for which the indices are
+      returned. If not specified, then ``m = n``.
+
+  Returns:
+    A tuple of two arrays containing the indices of the upper triangle, one along
+    each axis.
+
+  See also:
+    - :func:`jax.numpy.tril_indices`: Returns the indices of lower triangle of an
+      array of size ``(n, m)``.
+    - :func:`jax.numpy.triu_indices_from`: Returns the indices of upper triangle
+      of a given array.
+    - :func:`jax.numpy.tril_indices_from`: Returns the indices of lower triangle
+      of a given array.
+
+  Examples:
+    If only ``n`` is provided in input, the indices of upper triangle of an array
+    of size ``(n, n)`` array are returned.
+
+    >>> jnp.triu_indices(3)
+    (Array([0, 0, 0, 1, 1, 2], dtype=int32), Array([0, 1, 2, 1, 2, 2], dtype=int32))
+
+    If both ``n`` and ``m`` are provided in input, the indices of upper triangle
+    of an ``(n, m)`` array are returned.
+
+    >>> jnp.triu_indices(3, m=2)
+    (Array([0, 0, 1], dtype=int32), Array([0, 1, 1], dtype=int32))
+
+    If ``k = 1``, the indices on and above the first sub-diagonal above the main
+    diagonal are returned.
+
+    >>> jnp.triu_indices(3, k=1)
+    (Array([0, 0, 1], dtype=int32), Array([1, 2, 2], dtype=int32))
+
+    If ``k = -1``, the indices on and above the first sub-diagonal below the main
+    diagonal are returned.
+
+    >>> jnp.triu_indices(3, k=-1)
+    (Array([0, 0, 0, 1, 1, 1, 2, 2], dtype=int32), Array([0, 1, 2, 0, 1, 2, 1, 2], dtype=int32))
+  """
   n = core.concrete_or_error(operator.index, n, "n argument of jnp.triu_indices")
   k = core.concrete_or_error(operator.index, k, "k argument of jnp.triu_indices")
   m = n if m is None else core.concrete_or_error(operator.index, m, "m argument of jnp.triu_indices")
@@ -4962,8 +5011,57 @@ def triu_indices(n: int, k: int = 0, m: int | None = None) -> tuple[Array, Array
   return i, j
 
 
-@util.implements(np.tril_indices)
 def tril_indices(n: int, k: int = 0, m: int | None = None) -> tuple[Array, Array]:
+  """Return the indices of lower triangle of an array of size ``(n, m)``.
+
+  JAX implementation of :func:`numpy.tril_indices`.
+
+  Args:
+    n: int. Number of rows of the array for which the indices are returned.
+    k: optional, int, default=0. Specifies the sub-diagonal on and below which
+      the indices of lower triangle are returned. ``k=0`` refers to main diagonal,
+      ``k<0`` refers to sub-diagonal below the main diagonal and ``k>0`` refers
+      to sub-diagonal above the main diagonal.
+    m: optional, int. Number of columns of the array for which the indices are
+      returned. If not specified, then ``m = n``.
+
+  Returns:
+    A tuple of two arrays containing the indices of the lower triangle, one along
+    each axis.
+
+  See also:
+    - :func:`jax.numpy.triu_indices`: Returns the indices of upper triangle of an
+      array of size ``(n, m)``.
+    - :func:`jax.numpy.triu_indices_from`: Returns the indices of upper triangle
+      of a given array.
+    - :func:`jax.numpy.tril_indices_from`: Returns the indices of lower triangle
+      of a given array.
+
+  Examples:
+    If only ``n`` is provided in input, the indices of lower triangle of an array
+    of size ``(n, n)`` array are returned.
+
+    >>> jnp.tril_indices(3)
+    (Array([0, 1, 1, 2, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1, 2], dtype=int32))
+
+    If both ``n`` and ``m`` are provided in input, the indices of lower triangle
+    of an ``(n, m)`` array are returned.
+
+    >>> jnp.tril_indices(3, m=2)
+    (Array([0, 1, 1, 2, 2], dtype=int32), Array([0, 0, 1, 0, 1], dtype=int32))
+
+    If ``k = 1``, the indices on and below the first sub-diagonal above the main
+    diagonal are returned.
+
+    >>> jnp.tril_indices(3, k=1)
+    (Array([0, 0, 1, 1, 1, 2, 2, 2], dtype=int32), Array([0, 1, 0, 1, 2, 0, 1, 2], dtype=int32))
+
+    If ``k = -1``, the indices on and below the first sub-diagonal below the main
+    diagonal are returned.
+
+    >>> jnp.tril_indices(3, k=-1)
+    (Array([1, 2, 2], dtype=int32), Array([0, 0, 1], dtype=int32))
+  """
   n = core.concrete_or_error(operator.index, n, "n argument of jnp.triu_indices")
   k = core.concrete_or_error(operator.index, k, "k argument of jnp.triu_indices")
   m = n if m is None else core.concrete_or_error(operator.index, m, "m argument of jnp.triu_indices")
diff --git a/jax/_src/pallas/core.py b/jax/_src/pallas/core.py
index e589f9dda..eec1913b3 100644
--- a/jax/_src/pallas/core.py
+++ b/jax/_src/pallas/core.py
@@ -27,7 +27,6 @@ import warnings
 
 import jax
 from jax._src import api_util
-from jax._src import config
 from jax._src import core as jax_core
 from jax._src import deprecations
 from jax._src import linear_util as lu
@@ -41,8 +40,7 @@ import jax.numpy as jnp
 
 
 class DynamicGridDim:
-  def __repr__(self):
-    return "DynamicGridDim"
+  pass
 dynamic_grid_dim = DynamicGridDim()
 
 
@@ -175,27 +173,18 @@ def current_grid_env() -> GridEnv | None:
 
 
 class Mapped:
-  """Used as a block shape dimension to denote a mapped dimension.
-  A mapped dimension behaves like `1` except it is squeezed from the block.
-  See :ref:`pallas_blockspec` for more details.
-  """
-  def __repr__(self):
-    return "Mapped"
+  pass
 mapped = Mapped()
 
 
 @dataclasses.dataclass(frozen=True)
 class Unblocked:
   padding: tuple[tuple[int, int], ...] | None = None
-
-  def __repr__(self):
-    return f"Unblocked(padding={self.padding})"
 unblocked = Unblocked()
 
 
 class Blocked:
-  def __repr__(self):
-    return "Blocked"
+  pass
 blocked = Blocked()
 
 
@@ -207,8 +196,6 @@ class BlockSpec:
   """Specifies how an array should be sliced for each iteration of a kernel.
 
   See :ref:`pallas_blockspec` for more details.
-  This object contains the parameters passed through the API.
-  An internal canonicalized version is in BlockMapping.
   """
   block_shape: tuple[int | None, ...] | None = None
   index_map: Callable[..., Any] | None = None
@@ -260,38 +247,9 @@ BlockSpecTree = Any
 
 @dataclasses.dataclass(frozen=True)
 class BlockMapping:
-  """An internal canonicalized version of BlockSpec.
-
-  See the `check_invariants` method for precise specification.
-  """
   block_shape: tuple[Mapped | int, ...]
-  block_aval: AbstractMemoryRef   # The block ref aval
   index_map_jaxpr: jax_core.ClosedJaxpr
   indexing_mode: IndexingMode
-  array_shape_dtype: jax.ShapeDtypeStruct  # The whole array
-  origin: str  # The origin, e.g. input[2]["field"]
-
-  def check_invariants(self) -> None:
-    if not config.enable_checks.value: return
-
-    unmapped_block_shape = tuple(s for s in self.block_shape if s is not mapped)
-    assert unmapped_block_shape == self.block_aval.shape, (
-        self.block_shape, self.block_aval)
-    assert len(self.block_shape) == len(self.array_shape_dtype.shape), (
-        self.block_shape, self.array_shape_dtype
-    )
-
-    assert not self.index_map_jaxpr.consts
-    assert len(self.block_shape) == len(self.index_map_jaxpr.out_avals)
-    assert all(ov.shape == () and
-               (ov.dtype == jnp.int32 or ov.dtype == jnp.int64)
-               for ov in self.index_map_jaxpr.out_avals), (
-               self.index_map_jaxpr.out_avals)
-
-  def replace(self, **kwargs):
-    new_self = dataclasses.replace(self, **kwargs)
-    new_self.check_invariants()
-    return new_self
 
   def compute_start_indices_interpret(self, loop_idx, *args):
     discharged_jaxpr, discharged_consts = state_discharge.discharge_state(
@@ -311,6 +269,8 @@ class BlockMapping:
     else:
       raise RuntimeError(f"Unknown indexing mode: {self.indexing_mode}")
 
+  replace = dataclasses.replace
+
 
 @contextlib.contextmanager
 def tracing_grid_env(grid: GridMappingGrid, mapped_dims: tuple[int, ...]):
@@ -325,81 +285,16 @@ def tracing_grid_env(grid: GridMappingGrid, mapped_dims: tuple[int, ...]):
 
 @dataclasses.dataclass(frozen=True)
 class GridMapping:
-  """An internal canonicalized version of GridSpec.
-
-  Encodes the calling conventions of the pallas_call primitive, the kernel,
-  and the index maps.
-
-  The pallas_call is invoked with: ``*dynamic_grid_sizes, *index, *consts, *inputs``.
-  The ``index`` operands are for the scalar prefetch.
-  The ``consts`` are constants captured by the kernel function.
-
-  The kernel function is invoked with:
-  ``*index, *consts, *inputs, *scratch``.
-
-  The index map functions are invoked with:
-  ``*program_ids, *index``.
-
-  See the `check_invariants` method for a more precise specification.
-  """
   grid: GridMappingGrid
   grid_names: tuple[Hashable, ...] | None
+  block_mappings: tuple[BlockMapping | None, ...]
+  mapped_dims: tuple[int, ...] = ()
+  num_index_operands: int = 0
+  num_scratch_operands: int = 0
+  # Number of constants hoisted to operands by ``_hoist_consts_to_refs``.
+  num_constant_operands: int = 0
 
-  # Block mappings for: *consts, *inputs, *outputs
-  block_mappings: tuple[BlockMapping, ...]
-  # The inputs for tracing the index map: the tree and the flat avals
-  index_map_tree: tree_util.PyTreeDef
-  index_map_avals: tuple[jax_core.AbstractValue]
-  # Which dimensions in `grid` are vmapped.
-  vmapped_dims: tuple[int, ...]
-
-  num_index_operands: int
-  # Number of captured constants hoisted to operands.
-  num_constant_operands: int
-  num_inputs: int
-  num_outputs: int
-  num_scratch_operands: int
-
-  def check_invariants(self) -> None:
-    if not config.enable_checks.value: return
-    assert (len(self.block_mappings) ==
-            self.num_constant_operands + self.num_inputs + self.num_outputs), (
-        self.num_constant_operands, self.num_inputs, self.num_outputs,
-        self.block_mappings
-    )
-    # index_map_avals = int32[] * len(self.grid) + index_operands
-    assert len(self.index_map_avals) == len(self.grid) + self.num_index_operands, (
-        self.index_map_avals,
-        self.grid,
-        self.num_index_operands,
-    )
-    # Check that we can put together the avals and the tree.
-    index_map_args, index_map_kwargs = self.index_map_tree.unflatten(
-        self.index_map_avals)
-    assert not index_map_kwargs
-    assert len(index_map_args) >= len(self.grid)
-    for i in range(len(self.grid)):
-      index_map_arg = index_map_args[i]
-      assert index_map_arg.shape == ()
-      assert index_map_arg.dtype == jnp.int32
-
-    assert len(self.vmapped_dims) <= len(self.grid)
-    for i in self.vmapped_dims:
-      assert 0 <= i < len(self.grid)
-
-    if self.grid_names is not None:
-      assert len(self.grid) == len(self.grid_names), (self.grid, self.grid_names)
-    for bm in self.block_mappings:
-      bm.check_invariants()
-      assert tuple(self.index_map_avals) == tuple(bm.index_map_jaxpr.in_avals), (
-          self.index_map_avals,
-          bm.index_map_jaxpr.in_avals,
-      )
-
-  def replace(self, **kwargs) -> GridMapping:
-    new_self = dataclasses.replace(self, **kwargs)
-    new_self.check_invariants()
-    return new_self
+  replace = dataclasses.replace
 
   @property
   def num_dynamic_grid_bounds(self):
@@ -419,123 +314,74 @@ class GridMapping:
       axis_env_ctx = jax_core.extend_axis_env_nd(
           zip(self.grid_names, self.grid)
       )
-    with tracing_grid_env(self.grid, self.vmapped_dims), axis_env_ctx:
+    with tracing_grid_env(self.grid, self.mapped_dims), axis_env_ctx:
       yield
 
-  @property
-  def slice_index_ops(self):
-    """Returns a slice object to select the index operands to a kernel."""
-    return slice(0, self.num_index_operands)
-
-  @property
-  def slice_block_ops(self):
-    """Returns a slice to select all but the index operands to a kernel."""
-    return slice(self.num_index_operands, None)
-
-  @property
-  def slice_scratch_ops(self):
-    """Returns a slice object to select the scratch operands to a kernel."""
-    if self.num_scratch_operands:
-      return slice(-self.num_scratch_operands, None)
-    else:
-      return slice(0, 0)
-
-  # TODO(necula): this is used to recover the old `in_shapes`, but it probably
-  # is not needed anymore, with some cleanup.
-  @property
-  def in_shapes(self) -> tuple[jax.ShapeDtypeStruct, ...]:
-    """The shapes of *index, *consts, *inputs."""
-    index_shapes = [jax.ShapeDtypeStruct(ia.inner_aval.shape,
-                                         ia.inner_aval.dtype)
-                    for ia in self.index_map_avals[len(self.grid):]]
-    consts_inputs_shapes = [
-        bm.array_shape_dtype
-        for bm in self.block_mappings[
-            :self.num_constant_operands + self.num_inputs]]
-    return tuple(index_shapes + consts_inputs_shapes)
-
-  # TODO(necula): this is used to recover the old `out_shapes`, but it probably
-  # is not needed anymore, with some cleanup.
-  @property
-  def out_shapes(self) -> tuple[jax.ShapeDtypeStruct, ...]:
-    return tuple(
-        bm.array_shape_dtype
-        for bm in self.block_mappings[
-            self.num_constant_operands + self.num_inputs:
-            self.num_constant_operands + self.num_inputs + self.num_outputs])
-
 def _is_valid_grid_dim(dim: int | jax.Array) -> bool:
   if isinstance(dim, jax.Array):
     return True
   return jax_core.is_dim(dim)
 
+def _preprocess_grid(grid: Grid | int | None) -> tuple[TupleGrid, GridNames]:
+  if grid is None:
+    return (), None
+  if isinstance(grid, int):
+    return (grid,), None
+  # Handle empty grid
+  if not grid:
+    return grid, None  # type: ignore
+  # Check if we have a named grid
+  if isinstance(grid[0], tuple):
+    grid_names, grid = util.unzip2(grid)  # type: ignore
+  else:
+    grid_names = None
+  # TODO(b/353730556): allow NumPy scalars in grids
+  if not all(_is_valid_grid_dim(g) for g in grid):  # type: ignore
+    raise ValueError(
+        f"Grid must be a tuple of integers or jax.Array, got {grid}"
+    )
+  return grid, grid_names  # type: ignore
+
+
 def _convert_block_spec_to_block_mapping(
+    in_avals: Sequence[jax_core.ShapedArray],
     block_spec: BlockSpec,
     path: tree_util.KeyPath,
-    array_aval: jax_core.ShapedArray,
-    *,
-    # Inputs for the index_map
-    index_map_avals: Sequence[jax_core.AbstractValue],
-    index_map_tree: tree_util.PyTreeDef,
+    aval: jax_core.ShapedArray,
+    in_tree: Any,
     grid: GridMappingGrid,
     mapped_dims: tuple[int, ...],
     what: str,  # Used to localize error messages, e.g., {what}{path}
-) -> BlockMapping:
-  origin = f"{what}{tree_util.keystr(path)}"
+) -> BlockMapping | None:
   if block_spec is no_block_spec:
-    block_spec = BlockSpec(None, None)
+    return None
   if block_spec.index_map is None:
-    compute_index = lambda *args: (0,) * len(array_aval.shape)
+    compute_index = lambda *args, **kwargs: (0,) * len(aval.shape)
   else:
     compute_index = block_spec.compute_index
   if block_spec.block_shape is None:
-    block_shape = array_aval.shape
+    block_shape = aval.shape
   else:
     block_shape = block_spec.block_shape
-    if len(array_aval.shape) != len(block_shape):
+  block_shape = tuple(
+      mapped if s is None else s for s in block_shape)
+  flat_fun, _ = api_util.flatten_fun(lu.wrap_init(compute_index), in_tree)
+  with tracing_grid_env(grid, mapped_dims):
+    jaxpr, out_avals, consts, () = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals)
+    if len(out_avals) != len(block_shape):
       raise ValueError(
-          f"Block shape for {origin} (= {block_shape}) "
-          f"must have the same number of dimensions as the array shape {array_aval.shape}"
+          f"Index map for {what}{tree_util.keystr(path)} must return "
+          f"{len(aval.shape)} values to match {block_shape=}. "
+          f"Currently returning {len(out_avals)} values."
       )
-  unmapped_block_shape = tuple(s for s in block_shape if s is not None)
-  block_aval = AbstractMemoryRef(array_aval.update(shape=unmapped_block_shape),
-                                 block_spec.memory_space)
-
-  if not jax_core.is_constant_shape(block_aval.shape):
-    raise ValueError(
-        "shape polymorphism for Pallas does not support "
-        "dynamically-shaped blocks. "
-        f"{origin} has block_shape: {block_aval.shape}")
-
-  flat_index_map_fun, _ = api_util.flatten_fun(lu.wrap_init(compute_index),
-                                               index_map_tree)
-  with tracing_grid_env(grid, mapped_dims):
-    jaxpr, out_avals, consts, () = pe.trace_to_jaxpr_dynamic(flat_index_map_fun,
-                                                             index_map_avals)
-  mapped_block_shape = tuple(
-        mapped if s is None else s for s in block_shape)
-  if len(out_avals) != len(mapped_block_shape):
-    raise ValueError(
-        # TODO(necula): show the name and location of the index map function
-        f"Index map for {origin} must return "
-        f"{len(block_aval.shape)} values to match block shape {mapped_block_shape}. "
-        f"Currently returning {len(out_avals)} values."
-    )
   if consts:
     raise NotImplementedError(
-        # TODO(necula): show the name and location of the index map function
-        f"Index map for {origin} captures constants: "
+        f"Index map for {what}{tree_util.keystr(path)} captures constants: "
         f"{consts}")
-  mapping = BlockMapping(
-      block_shape=mapped_block_shape,
-      block_aval=block_aval,
-      index_map_jaxpr=jax_core.ClosedJaxpr(jaxpr, consts),
-      indexing_mode=block_spec.indexing_mode,
-      array_shape_dtype=jax.ShapeDtypeStruct(array_aval.shape, array_aval.dtype),
-      origin=origin,
+  return BlockMapping(
+      block_shape, jax_core.ClosedJaxpr(jaxpr, consts), block_spec.indexing_mode
   )
-  mapping.check_invariants()
-  return mapping
+
 
 def _tile_ref(ref: state.AbstractRef, block_shape: tuple[int, ...] | None
              ) -> state.AbstractRef:
@@ -544,20 +390,65 @@ def _tile_ref(ref: state.AbstractRef, block_shape: tuple[int, ...] | None
   shape = tuple(s for s in block_shape if s is not None)
   return ref.update(inner_aval=ref.inner_aval.update(shape=shape))
 
+
+def _get_ref_avals(in_avals: Sequence[jax_core.ShapedArray],
+                   in_specs: Sequence[BlockSpec],
+                   in_paths: Sequence[tree_util.KeyPath],
+                   out_avals: Sequence[jax_core.ShapedArray],
+                   out_specs: Sequence[BlockSpec],
+                   out_paths: Sequence[tree_util.KeyPath]):
+  def make_ref_aval(aval: jax_core.ShapedArray,
+                    spec: BlockSpec,
+                    path: tree_util.KeyPath,
+                    what: str) -> state.AbstractRef:
+    if spec is no_block_spec:
+      memory_space = None
+      block_shape = None
+    else:
+      memory_space = spec.memory_space
+      block_shape = spec.block_shape
+
+    ref_aval = AbstractMemoryRef(aval, memory_space)
+    if block_shape is not None:
+      if len(ref_aval.shape) != len(block_shape):
+        raise ValueError(
+            f"Block shape for {what}{tree_util.keystr(path)} (= {block_shape}) "
+            f"must have the same number of dimensions as the array shape {ref_aval.shape}"
+        )
+      block_shape_unmapped = tuple(s for s in block_shape if s is not None)
+      ref_aval = ref_aval.update(
+          inner_aval=ref_aval.inner_aval.update(shape=block_shape_unmapped))
+
+    if not jax_core.is_constant_shape(ref_aval.shape):
+      raise ValueError(
+          "shape polymorphism for Pallas does not support "
+          "dynamically-shaped blocks. "
+          f"{what}{tree_util.keystr(path)} has block_shape: {ref_aval.shape}")
+    return ref_aval
+
+  in_ref_avals = [
+      make_ref_aval(aval, in_spec, in_path, "input")
+      for aval, in_spec, in_path in zip(in_avals, in_specs, in_paths)
+  ]
+  out_ref_avals = [
+      make_ref_aval(aval, out_spec, out_path, "output")
+      for aval, out_spec, out_path in zip(out_avals, out_specs, out_paths)
+  ]
+  return in_ref_avals, out_ref_avals
+
+
 @dataclasses.dataclass(init=False, unsafe_hash=True)
 class GridSpec:
-  """Encodes the parameters of the grid, as given through the API.
-
-  An internal sanitized version is in GridMapping.
-  """
   grid: TupleGrid
   grid_names: tuple[Hashable, ...] | None
-  in_specs: tuple[BlockSpec | NoBlockSpec, ...] | NoBlockSpec
-  out_specs: tuple[BlockSpec | NoBlockSpec, ...] | NoBlockSpec
+  in_specs: tuple[BlockSpec | NoBlockSpec, ...]
+  out_specs: tuple[BlockSpec | NoBlockSpec, ...]
+  in_specs_tree: Any
+  out_specs_tree: Any
 
   def __init__(
       self,
-      grid: Grid = (),
+      grid: Grid | None = None,
       in_specs: BlockSpecTree = no_block_spec,
       out_specs: BlockSpecTree = no_block_spec,
   ):
@@ -569,155 +460,89 @@ class GridSpec:
     if isinstance(out_specs, list):
       out_specs = tuple(out_specs)
 
-    self.in_specs = in_specs
-    self.out_specs = out_specs
-
-    grid_names = None
-    if isinstance(grid, int):
-      grid = (grid,)
-    elif grid and isinstance(grid[0], tuple):  # Check if we have a named grid
-      grid_names, grid = util.unzip2(grid)  # type: ignore
-
-    # TODO(b/353730556): allow NumPy scalars in grids
-    if not all(_is_valid_grid_dim(g) for g in grid):  # type: ignore
-      raise ValueError(
-          f"Grid must be a tuple of integers or jax.Array, got {grid}"
-      )
-    self.grid = grid  # type: ignore
-    self.grid_names = grid_names
-
-  def get_grid_mapping(
-      self,
-      in_avals: Sequence[jax_core.AbstractValue],
-      in_tree: tree_util.PyTreeDef,
-      in_paths: Sequence[tree_util.KeyPath],
-      out_avals: Sequence[jax_core.AbstractValue],
-      out_tree: tree_util.PyTreeDef,
-      out_paths: Sequence[tree_util.KeyPath],
-      grid_names: tuple[Hashable, ...] | None = None,
-      num_scalar_prefetch: int = 0,
-      scratch_shapes: Sequence[Any] = (),
-  ) -> tuple[tuple[jax_core.AbstractValue, ...],
-             GridMapping]:
-    assert all(i is None or isinstance(i, int) for i in self.grid)
-    grid_mapping_grid = tuple(
-        dynamic_grid_dim if d is None else d for d in self.grid
-    )
-    # The inputs for the index maps
-    index_map_avals = (
-        (jax_core.ShapedArray((), jnp.dtype("int32")),) * len(self.grid))
-    index_map_tree = tree_util.tree_structure((index_map_avals, {}))
-
-    if num_scalar_prefetch:
-      all_avals = tree_util.tree_unflatten(in_tree, in_avals)
-      scalar_avals, unflat_in_avals = split_list(
-          all_avals, [num_scalar_prefetch])
-      flat_scalar_avals, scalar_tree = tree_util.tree_flatten(scalar_avals)
-      num_flat_scalar_prefetch = len(flat_scalar_avals)
-      scalar_ref_avals = [
-          self._make_scalar_ref_aval(aval)
-          for aval in flat_scalar_avals]
-      jaxpr_scalar_ref_avals = tree_util.tree_unflatten(
-          scalar_tree, scalar_ref_avals)
-      in_avals, in_tree = tree_util.tree_flatten(tuple(unflat_in_avals))
-      index_map_tree = tree_util.tree_structure(((*index_map_avals,
-                                                  *scalar_avals), {}))
-      index_map_avals = (*index_map_avals, *scalar_ref_avals)
-      del scalar_ref_avals, flat_scalar_avals, scalar_tree
-      del scalar_avals, unflat_in_avals, all_avals
+    self.grid, self.grid_names = _preprocess_grid(grid)
+    if in_specs is not no_block_spec:
+      flat_in_specs, self.in_specs_tree = tree_util.tree_flatten(in_specs)
+      self.in_specs = tuple(flat_in_specs)
     else:
-      num_flat_scalar_prefetch = 0
-      jaxpr_scalar_ref_avals = ()
-
-    if scratch_shapes:
-      flat_scratch_shapes, scratch_tree = tree_util.tree_flatten(
-          scratch_shapes)
-      flat_scratch_avals = map(self._make_scratch_aval, flat_scratch_shapes)
-      num_flat_scratch_operands = len(flat_scratch_avals)
-      jaxpr_scratch_avals = tree_util.tree_unflatten(
-          scratch_tree, flat_scratch_avals)
-      if not isinstance(jaxpr_scratch_avals, (tuple, list)):
-        jaxpr_scratch_avals = (jaxpr_scratch_avals,)
-      del flat_scratch_avals, flat_scratch_shapes, scratch_tree
+      self.in_specs = in_specs
+      self.in_specs_tree = None
+    if out_specs is not no_block_spec:
+      flat_out_specs, self.out_specs_tree = tree_util.tree_flatten(out_specs)
+      self.out_specs = tuple(flat_out_specs)
     else:
-      num_flat_scratch_operands = 0
-      jaxpr_scratch_avals = ()
+      self.out_specs = out_specs
+      self.out_specs_tree = None
 
-    if self.in_specs is not no_block_spec:
-      flat_in_specs, in_specs_tree = tree_util.tree_flatten(self.in_specs)
-      if in_specs_tree != in_tree:
+  def _get_in_out_specs(self, in_avals, in_tree, out_avals, out_tree):
+    if self.in_specs is no_block_spec:
+      flat_in_specs = [no_block_spec] * len(in_avals)
+    else:
+      flat_in_specs = self.in_specs
+      if self.in_specs_tree != in_tree:
         raise ValueError(
-            pytreedef_mismatch_err_msg("`in_specs`", in_specs_tree,
+            pytreedef_mismatch_err_msg("`in_specs`", self.in_specs_tree,
                                        "inputs", in_tree))
+    if self.out_specs is no_block_spec:
+      flat_out_specs = [no_block_spec] * len(out_avals)
     else:
-      flat_in_specs = [no_block_spec] * len(in_avals)
+      flat_out_specs = self.out_specs
+      if self.out_specs_tree != out_tree:
+        raise ValueError(
+            pytreedef_mismatch_err_msg("`out_specs`", self.out_specs_tree,
+                                       "`out_shape`", out_tree))
+    return flat_in_specs, flat_out_specs
 
+  def get_grid_mapping(
+      self, in_avals, in_tree, in_paths, out_avals, out_tree, out_paths
+  ) -> tuple[tuple[jax_core.AbstractValue, ...], GridMapping]:
+    assert all(i is None or isinstance(i, int) for i in self.grid)
+    grid_mapping_grid = tuple(
+        dynamic_grid_dim if d is None else d for d in self.grid
+    )
+    flat_in_specs, flat_out_specs = self._get_in_out_specs(
+        in_avals, in_tree, out_avals, out_tree)
+    in_ref_avals, out_ref_avals = _get_ref_avals(
+        in_avals, flat_in_specs, in_paths,
+        out_avals, flat_out_specs, out_paths)
+    grid_avals = [jax_core.ShapedArray((), jnp.dtype("int32"))] * len(self.grid)
+    # Create args, kwargs pytree def
+    grid_tree = tree_util.tree_structure((tuple(grid_avals), {}))
     in_block_mappings = map(
         partial(
             _convert_block_spec_to_block_mapping,
-            index_map_avals=index_map_avals,
-            index_map_tree=index_map_tree,
+            grid_avals,
+            in_tree=grid_tree,
             grid=grid_mapping_grid,
             mapped_dims=(),
-            what="inputs",
+            what="input",
         ),
         flat_in_specs,
-        in_paths[num_flat_scalar_prefetch:],
-        in_avals,
+        in_paths,
+        in_ref_avals,
     )
-
-    if self.out_specs is not no_block_spec:
-      flat_out_specs, out_specs_tree = tree_util.tree_flatten(self.out_specs)
-      if out_specs_tree != out_tree:
-        raise ValueError(
-            pytreedef_mismatch_err_msg("`out_specs`", out_specs_tree,
-                                       "`out_shape`", out_tree))
-    else:
-      flat_out_specs = [no_block_spec] * len(out_avals)
-
     out_block_mappings = map(
         partial(
             _convert_block_spec_to_block_mapping,
-            index_map_avals=index_map_avals,
-            index_map_tree=index_map_tree,
+            grid_avals,
+            in_tree=grid_tree,
             grid=grid_mapping_grid,
             mapped_dims=(),
-            what="outputs",
+            what="output",
         ),
         flat_out_specs,
         out_paths,
-        out_avals,
+        out_ref_avals,
     )
     grid_mapping = GridMapping(
-        grid=grid_mapping_grid,  # type: ignore[arg-type]
-        grid_names=self.grid_names,
-        block_mappings=(*in_block_mappings, *out_block_mappings),
-        index_map_avals=index_map_avals,  # type: ignore[arg-type]
-        index_map_tree=index_map_tree,
-        vmapped_dims=(),
-        num_index_operands=num_flat_scalar_prefetch,
-        num_constant_operands=0,  # Fixed up later
-        num_inputs=len(flat_in_specs),
-        num_outputs=len(flat_out_specs),
-        num_scratch_operands=num_flat_scratch_operands,
+        grid_mapping_grid, self.grid_names, # type: ignore
+        (*in_block_mappings, *out_block_mappings)
     )
-    grid_mapping.check_invariants()
-    in_ref_avals = [bm.block_aval for bm in in_block_mappings]
-    jaxpr_in_ref_avals = tree_util.tree_unflatten(in_tree, in_ref_avals)
-    jaxpr_in_avals = (*jaxpr_scalar_ref_avals,
-                      *jaxpr_in_ref_avals)
-    out_ref_avals = [bm.block_aval for bm in out_block_mappings]
+    jaxpr_in_avals = tree_util.tree_unflatten(in_tree, in_ref_avals)
     jaxpr_out_avals = tree_util.tree_unflatten(out_tree, out_ref_avals)
     if not isinstance(jaxpr_out_avals, (tuple, list)):
       jaxpr_out_avals = (jaxpr_out_avals,)
-    return (*jaxpr_in_avals, *jaxpr_out_avals,
-            *jaxpr_scratch_avals), grid_mapping
-
-  def _make_scratch_aval(self, obj: object) -> jax_core.AbstractValue:
-    assert False  # Not needed in GridSpec
-
-  def _make_scalar_ref_aval(self, aval):
-    assert False  # Not needed in GridSpec
+    return (*jaxpr_in_avals, *jaxpr_out_avals), grid_mapping
 
   def unzip_dynamic_grid_bounds(
       self,
@@ -732,7 +557,6 @@ class GridSpec:
     static_self.grid = static_grid  # type: ignore
     return static_self, dynamic_bounds
 
-
 def pytreedef_mismatch_err_msg(
     what1: str, tree1: tree_util.PyTreeDef,
     what2: str, tree2: tree_util.PyTreeDef) -> str:
diff --git a/jax/_src/pallas/mosaic/core.py b/jax/_src/pallas/mosaic/core.py
index 59461be81..d5e62ef1f 100644
--- a/jax/_src/pallas/mosaic/core.py
+++ b/jax/_src/pallas/mosaic/core.py
@@ -24,6 +24,7 @@ from typing import Any, Hashable
 import jax
 from jax._src import core as jax_core
 from jax._src import dtypes
+from jax._src import tree_util
 from jax._src import util
 import jax.numpy as jnp
 from jax._src.pallas import core as pallas_core
@@ -140,19 +141,30 @@ class MemoryRef:
         jax_core.ShapedArray(self.shape, self.dtype), self.memory_space)
 
 
+def _make_aval(obj: object) -> jax_core.AbstractValue:
+  if isinstance(obj, MemoryRef):
+    return obj.get_aval()
+  if isinstance(obj, SemaphoreType):
+    return obj.get_aval()
+  raise ValueError(f"No registered conversion for {type(obj)}. "
+                   "Only VMEM and SemaphoreType are supported.")
+
+
 @dataclasses.dataclass(init=False, unsafe_hash=True)
 class PrefetchScalarGridSpec(pallas_core.GridSpec):
   grid: TupleGrid
   grid_names: tuple[Hashable, ...] | None
   num_scalar_prefetch: int
-  in_specs: tuple[BlockSpec | NoBlockSpec, ...] | NoBlockSpec
-  out_specs: tuple[BlockSpec | NoBlockSpec, ...] | NoBlockSpec
+  in_specs: tuple[BlockSpec | NoBlockSpec, ...]
+  out_specs: tuple[BlockSpec | NoBlockSpec, ...]
+  in_specs_tree: Any
+  out_specs_tree: Any
   scratch_shapes: tuple[Any, ...]
 
   def __init__(
       self,
       num_scalar_prefetch: int,
-      grid: Grid = (),
+      grid: Grid | None = None,
       in_specs: BlockSpecTree = no_block_spec,
       out_specs: BlockSpecTree = no_block_spec,
       scratch_shapes: Any | Sequence[Any] = ()
@@ -161,26 +173,84 @@ class PrefetchScalarGridSpec(pallas_core.GridSpec):
     self.num_scalar_prefetch = num_scalar_prefetch
     self.scratch_shapes = tuple(scratch_shapes)
 
-  def _make_scalar_ref_aval(self, aval):
-    return AbstractMemoryRef(jax_core.ShapedArray(aval.shape, aval.dtype),
-                             TPUMemorySpace.SMEM)
-
-  def _make_scratch_aval(self, obj: object) -> jax_core.AbstractValue:
-    if isinstance(obj, MemoryRef):
-      return obj.get_aval()
-    if isinstance(obj, SemaphoreType):
-      return obj.get_aval()
-    raise ValueError(f"No registered conversion for {type(obj)}. "
-                     "Only VMEM and SemaphoreType are supported.")
-
-  def get_grid_mapping(  # type: ignore[override]
+  def get_grid_mapping(
       self, in_avals, in_tree, in_paths, out_avals, out_tree, out_paths
   ) -> tuple[tuple[jax_core.AbstractValue, ...], GridMapping]:
-    return super().get_grid_mapping(in_avals, in_tree, in_paths,
-                                    out_avals, out_tree, out_paths,
-                                    grid_names=self.grid_names,
-                                    num_scalar_prefetch=self.num_scalar_prefetch,
-                                    scratch_shapes=self.scratch_shapes)
+    assert all(i is None or isinstance(i, int) for i in self.grid)
+    grid_mapping_grid = tuple(
+        pallas_core.dynamic_grid_dim if d is None else d for d in self.grid
+    )
+    all_avals = tree_util.tree_unflatten(in_tree, in_avals)
+    flat_scratch_shapes, scratch_tree = tree_util.tree_flatten(
+        self.scratch_shapes)
+    flat_scratch_avals = map(_make_aval, flat_scratch_shapes)
+    scalar_avals, unflat_in_avals = split_list(
+        all_avals, [self.num_scalar_prefetch])
+    flat_scalar_avals, scalar_tree = tree_util.tree_flatten(scalar_avals)
+    num_flat_scalar_prefetch = len(flat_scalar_avals)
+    in_avals, in_avals_tree = tree_util.tree_flatten(tuple(unflat_in_avals))
+    flat_in_specs, flat_out_specs = self._get_in_out_specs(
+        in_avals, in_avals_tree, out_avals, out_tree)
+    in_ref_avals, out_ref_avals = (
+        pallas_core._get_ref_avals(
+            in_avals, flat_in_specs, in_paths[num_flat_scalar_prefetch:],
+            out_avals, flat_out_specs, out_paths))
+    scalar_ref_avals = [
+        AbstractMemoryRef(jax_core.ShapedArray(aval.shape, aval.dtype),
+                          TPUMemorySpace.SMEM)
+        for aval in flat_scalar_avals]
+    grid_avals = [jax_core.ShapedArray((), jnp.dtype("int32"))] * len(self.grid)
+    # Create args, kwargs pytree def
+    index_map_in_tree = tree_util.tree_structure(
+        ((*grid_avals, *scalar_avals), {})
+    )
+    in_block_mappings = map(
+        partial(
+            _convert_block_spec_to_block_mapping,
+            (*grid_avals, *scalar_ref_avals),
+            in_tree=index_map_in_tree,
+            grid=grid_mapping_grid,
+            mapped_dims=(),
+            what="input",
+        ),
+        flat_in_specs,
+        in_paths[num_flat_scalar_prefetch:],
+        in_ref_avals,
+    )
+    out_block_mappings = map(
+        partial(
+            _convert_block_spec_to_block_mapping,
+            (*grid_avals, *scalar_ref_avals),
+            in_tree=index_map_in_tree,
+            grid=grid_mapping_grid,
+            mapped_dims=(),
+            what="output",
+        ),
+        flat_out_specs,
+        out_paths,
+        out_ref_avals,
+    )
+    grid_mapping = GridMapping(
+        grid=grid_mapping_grid, grid_names=self.grid_names, # type: ignore
+        block_mappings=(*in_block_mappings, *out_block_mappings),
+        mapped_dims=(),
+        num_index_operands=num_flat_scalar_prefetch,
+        num_scratch_operands=len(flat_scratch_avals)
+    )
+    jaxpr_scalar_ref_avals = tree_util.tree_unflatten(
+        scalar_tree, scalar_ref_avals)
+    jaxpr_in_ref_avals = tree_util.tree_unflatten(in_avals_tree, in_ref_avals)
+    jaxpr_scratch_avals = tree_util.tree_unflatten(
+        scratch_tree, flat_scratch_avals)
+    if not isinstance(jaxpr_scratch_avals, (tuple, list)):
+      jaxpr_scratch_avals = (jaxpr_scratch_avals,)
+    jaxpr_in_avals = (*jaxpr_scalar_ref_avals,
+                      *jaxpr_in_ref_avals)
+    jaxpr_out_avals = tree_util.tree_unflatten(out_tree, out_ref_avals)
+    if not isinstance(jaxpr_out_avals, (tuple, list)):
+      jaxpr_out_avals = (jaxpr_out_avals,)
+    return (*jaxpr_in_avals, *jaxpr_out_avals,
+            *jaxpr_scratch_avals), grid_mapping
 
 
 @dataclasses.dataclass(frozen=True)
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index d67ad6838..b8905916a 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -294,7 +294,7 @@ class MosaicGridMapping:
     self.grid_names = grid_mapping.grid_names
     self.jaxpr = jaxpr
     self.block_mappings = grid_mapping.block_mappings
-    self.mapped_dims = grid_mapping.vmapped_dims
+    self.mapped_dims = grid_mapping.mapped_dims
     num_scalar_prefetch = grid_mapping.num_index_operands
     num_scratch = grid_mapping.num_scratch_operands
     # jaxpr has signature [*scalar_prefetch, *consts, *in_ops, *out_ops, *scratch]
@@ -425,15 +425,13 @@ class MeshInfo:
 def lower_jaxpr_to_module(
     ctx: ir.Context,
     grid_mapping: pl_core.GridMapping,
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     jaxpr: jax_core.Jaxpr,
     dimension_semantics: tuple[str | None, ...] | None,
     mesh: mesh_lib.Mesh | None = None,
     for_verification: bool = False,
 ) -> tuple[Module, tuple[Any, ...]]:
-  # TODO(necula): cleanup
-  in_shapes = grid_mapping.in_shapes
-  out_shapes = grid_mapping.out_shapes
-
   mosaic_grid_mapping = MosaicGridMapping(
       jaxpr, grid_mapping, dimension_semantics, mesh)
   mosaic_grid_mapping.maybe_compress_grid()
@@ -456,8 +454,10 @@ def lower_jaxpr_to_module(
       invars = invars[grid_mapping.num_index_operands:]
     # invars now = *consts, *ins, *outs
     avals = tuple(v.aval for v in invars)
-    # TODO(necula): we should not need block_operand_shapes anymore
     block_operand_shapes = (
+        *[jax.ShapeDtypeStruct(v.aval.shape,
+                               v.aval.dtype)
+          for v in invars[:grid_mapping.num_constant_operands]],
         *in_shapes[grid_mapping.num_index_operands:],
         *out_shapes,
     )
@@ -466,6 +466,10 @@ def lower_jaxpr_to_module(
         zip(block_operand_shapes, grid_mapping.block_mappings, avals)
     ):
       func_name = f"transform_{i}"
+      if bm is None:
+        raise NotImplementedError(
+            "BlockSpecs are required on TPU when grid is specified"
+        )
       # ANY operands don't support windowing and require empty window_params.
       if aval.memory_space == tpu_core.TPUMemorySpace.ANY:
         # We may not require windowing if our block_shape matches the original
diff --git a/jax/_src/pallas/mosaic/pallas_call_registration.py b/jax/_src/pallas/mosaic/pallas_call_registration.py
index 951a3f158..9e2486eb1 100644
--- a/jax/_src/pallas/mosaic/pallas_call_registration.py
+++ b/jax/_src/pallas/mosaic/pallas_call_registration.py
@@ -32,6 +32,7 @@ from jax._src.lib.mlir import ir
 from jax._src.pallas import core
 from jax._src.pallas.mosaic import lowering
 from jax._src.pallas.mosaic import verification
+from jax._src.pallas.pallas_call import pallas_call_p
 from jax.experimental import mosaic
 from jax.experimental.mosaic.dialects import tpu
 
@@ -68,13 +69,21 @@ def pallas_call_tpu_lowering_rule(
     name: str,
     grid_mapping: core.GridMapping,
     input_output_aliases: tuple[tuple[int, int], ...],
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     debug: bool,
     interpret: bool,
     compiler_params: dict[str, Any]):
   """Lowers a pallas_call to a Mosaic TPU custom call."""
-  del interpret
-  # TODO(necula): cleanup
-  out_shapes = grid_mapping.out_shapes
+  if interpret:
+    # TODO(necula): is this branch still needed?
+    return mlir.lower_fun(pallas_call_p.impl, multiple_results=True)(
+        ctx, *in_nodes, jaxpr=jaxpr, name=name, out_shapes=out_shapes,
+        in_shapes=in_shapes,
+        interpret=interpret, debug=debug,
+        input_output_aliases=input_output_aliases,
+        grid_mapping=grid_mapping,
+        compiler_params=compiler_params)
   if debug:
     print(jaxpr)
   if "mosaic_params" in compiler_params:
@@ -105,7 +114,7 @@ def pallas_call_tpu_lowering_rule(
     with mlir_ctx, ir.Location.unknown(mlir_ctx):
       dimension_semantics = mosaic_params.get("dimension_semantics", None)
       return lowering.lower_jaxpr_to_module(
-          mlir_ctx, grid_mapping, jaxpr,
+          mlir_ctx, grid_mapping, in_shapes, out_shapes, jaxpr,
           dimension_semantics=dimension_semantics, mesh=mesh,
           for_verification=for_verification)
   mosaic_module, extra_args = lower_module(for_verification=False)
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index d90a92ab0..c1f07d52b 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -156,14 +156,14 @@ class LoweringError(Exception):
 
 def lower_jaxpr_to_module(
     grid_mapping: pl_core.GridMapping,
+    in_structs: tuple[jax.ShapeDtypeStruct, ...],
+    out_structs: tuple[jax.ShapeDtypeStruct, ...],
     jaxpr: jax_core.Jaxpr,
     name: str,
     compiler_params: dict[str, Any],
 ) -> LoweringResult:
-  in_structs = grid_mapping.in_shapes
-  out_structs = grid_mapping.out_shapes
   assert len(jaxpr.outvars) == 0
-  assert not grid_mapping.vmapped_dims
+  assert not grid_mapping.mapped_dims
   grid = grid_mapping.grid
   if len(grid) < 3:
     grid += (1,) * (3 - len(grid))
diff --git a/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py b/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py
index 2f117625b..6aca2b74e 100644
--- a/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py
+++ b/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py
@@ -19,10 +19,12 @@ from __future__ import annotations
 
 from typing import Any
 
+import jax
 from jax import core as jax_core
 from jax._src.interpreters import mlir
 from jax._src.pallas import core as pallas_core
 from jax._src.pallas.mosaic_gpu import lowering
+from jax._src.pallas.pallas_call import pallas_call_p
 from jax.experimental.mosaic import gpu as mosaic_gpu
 
 
@@ -31,13 +33,30 @@ def pallas_call_lowering(
     *args,
     jaxpr: jax_core.Jaxpr,
     name: str,
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     interpret: bool,
     debug: bool,
     input_output_aliases: tuple[tuple[int, int], ...],
     grid_mapping: pallas_core.GridMapping,
     compiler_params: dict[str, Any],
 ):
-  del interpret
+  if interpret:
+    # TODO(necula): is this still needed?
+    return mlir.lower_fun(pallas_call_p.impl, multiple_results=True)(
+        ctx,
+        *args,
+        jaxpr=jaxpr,
+        name=name,
+        out_shapes=out_shapes,
+        in_shapes=in_shapes,
+        interpret=interpret,
+        debug=debug,
+        input_output_aliases=input_output_aliases,
+        grid_mapping=grid_mapping,
+        compiler_params=compiler_params,
+    )
+
   if grid_mapping.num_dynamic_grid_bounds:
     raise NotImplementedError(
         "dynamic grid bounds not supported in the Mosaic GPU backend"
@@ -56,6 +75,8 @@ def pallas_call_lowering(
     )
   lowering_result = lowering.lower_jaxpr_to_module(
       grid_mapping,
+      in_shapes,
+      out_shapes,
       jaxpr,
       name,
       compiler_params,
diff --git a/jax/_src/pallas/pallas_call.py b/jax/_src/pallas/pallas_call.py
index a950890fa..97770ff1a 100644
--- a/jax/_src/pallas/pallas_call.py
+++ b/jax/_src/pallas/pallas_call.py
@@ -61,7 +61,6 @@ BlockSpecTree = pallas_core.BlockSpecTree
 NoBlockSpec = pallas_core.NoBlockSpec
 no_block_spec = pallas_core.no_block_spec
 
-# See the docstring for GridMapping for the calling convention
 pallas_call_p = jax_core.Primitive('pallas_call')
 pallas_call_p.multiple_results = True
 
@@ -104,6 +103,8 @@ def _pad_values_to_block_dimension(value,
   Returns:
     A padded array.
   """
+  if block_shape is None:
+    return value
   padded_shape = tuple(
       ((v - 1) // b + 1) * b for v, b in zip(value.shape, block_shape)
   )
@@ -167,14 +168,12 @@ def _pallas_call_impl(*args, **kwargs):
 def _pallas_call_impl_interpret(
     *args,
     jaxpr: jax_core.Jaxpr,
-    name: str,
+    name: str, in_shapes, out_shapes,
     debug: bool,
     input_output_aliases: tuple[tuple[int, int], ...],
     grid_mapping: GridMapping,
     compiler_params: Any):
-  del compiler_params, name
-  # TODO(necula): cleanup
-  out_shapes = grid_mapping.out_shapes
+  del compiler_params, name, in_shapes
   # If we're in interpreter mode, we *scan* over the grid and eval the
   # discharged jaxpr.
   dynamic_grid_args, args = split_list(  # type: ignore
@@ -195,13 +194,21 @@ def _pallas_call_impl_interpret(
   scalars, args = split_list(args, [grid_mapping.num_index_operands])  # type: ignore
   # invars: [*scalar_prefetch, *consts, *inputs, *outputs, *scratch]
   # args now contains: *consts, *inputs, *outputs
-  scratch_invars = jaxpr.invars[grid_mapping.slice_scratch_ops]
+  num_invars = len(jaxpr.invars)
+  num_inputs_outputs = (
+      num_invars
+      - grid_mapping.num_index_operands
+      - grid_mapping.num_scratch_operands
+  )
+  _, _, scratch_invars = split_list(
+      jaxpr.invars, [grid_mapping.num_index_operands, num_inputs_outputs]
+  )
   scratch_avals = [v.aval for v in scratch_invars]
   scratch_values = _initialize_scratch_vals(scratch_avals)
 
   carry = []
   for x, bm in zip(itertools.chain(args, out), grid_mapping.block_mappings):
-    if isinstance(bm.indexing_mode, pallas_core.Unblocked):
+    if bm is not None and isinstance(bm.indexing_mode, pallas_core.Unblocked):
       padding = bm.indexing_mode.padding
       if padding is not None and any(p != (0, 0) for p in padding):
         if input_output_aliases:
@@ -210,14 +217,18 @@ def _pallas_call_impl_interpret(
         x = lax.pad(x, pad_value, [(*p, 0) for p in padding])
     carry.append(x)
 
+  block_shapes_without_mapped_dims = [
+      None if block_mapping is None else block_mapping.block_shape
+      for block_mapping in grid_mapping.block_mappings
+  ]
   is_indexing_dim = [
-      tuple(b is pallas_core.mapped for b in bm.block_shape)
-      for bm in grid_mapping.block_mappings
+      None if bm is None else tuple(b is pallas_core.mapped for b in bm)
+      for bm in block_shapes_without_mapped_dims
   ]
   block_shapes = [
-      None if iid is None
-      else tuple(1 if i else b for i, b in zip(iid, bm.block_shape))
-      for iid, bm in zip(is_indexing_dim, grid_mapping.block_mappings)
+      None if (bm is None or iid is None)
+      else tuple(1 if i else b for i, b in zip(iid, bm))
+      for iid, bm in zip(is_indexing_dim, block_shapes_without_mapped_dims)
   ]
 
   # Pad values to evenly divide into block dimensions. This matches the
@@ -245,7 +256,7 @@ def _pallas_call_impl_interpret(
     local_grid_env = tuple(
         pallas_core.GridAxis(idx, b)
         for dim, (idx, b) in enumerate(zip(loop_idx, grid))
-        if dim not in grid_mapping.vmapped_dims
+        if dim not in grid_mapping.mapped_dims
     )
     carry, scratch = split_list(carry, [num_inout])
     with pallas_core.grid_env(local_grid_env):
@@ -279,7 +290,7 @@ def _pallas_call_impl_interpret(
   out_block_mappings = grid_mapping.block_mappings[len(args):]
   out_nopad = []
   for o, expected_o_shape, bm in zip(out, out_shapes, out_block_mappings):
-    if isinstance(bm.indexing_mode, pallas_core.Unblocked):
+    if bm is not None and isinstance(bm.indexing_mode, pallas_core.Unblocked):
       padding = bm.indexing_mode.padding
       if padding is not None and any(p != (0, 0) for p in padding):
         if input_output_aliases:
@@ -294,16 +305,13 @@ def _pallas_call_impl_interpret(
 
 pallas_call_p.def_impl(_pallas_call_impl)
 
-def _pallas_call_abstract_eval(*avals, grid_mapping, **_):
-  out_shapes = grid_mapping.out_shapes
+def _pallas_call_abstract_eval(*avals, out_shapes, **_):
   return map(lambda x: jax_core.ShapedArray(x.shape, x.dtype), out_shapes)
 pallas_call_p.def_abstract_eval(_pallas_call_abstract_eval)
 
 def _pallas_call_jvp_rule(primals, tangents, *, jaxpr, name,
     input_output_aliases: tuple[tuple[int, int], ...],
-    grid_mapping, debug, interpret, compiler_params: Any):
-  # TODO(necula): cleanup
-  out_shapes = grid_mapping.out_shapes
+    in_shapes, out_shapes, grid_mapping, debug, interpret, compiler_params: Any):
   if grid_mapping.num_dynamic_grid_bounds:
     raise NotImplementedError("interpret with dynamic grid bounds unsupported")
   if grid_mapping.num_index_operands:
@@ -339,17 +347,14 @@ def _pallas_call_jvp_rule(primals, tangents, *, jaxpr, name,
     print(jvp_jaxpr)
   in_bms, out_bms = split_list(grid_mapping.block_mappings, [len(primals)])
   jvp_bms = (*in_bms, *in_bms, *out_bms, *out_bms)
-  jvp_grid_mapping = grid_mapping.replace(
-      block_mappings=jvp_bms,
-      num_inputs=grid_mapping.num_inputs * 2,
-      num_outputs=grid_mapping.num_outputs * 2,
-  )
   out_flat = pallas_call_p.bind(
       *primals,
       *tangents,
       jaxpr=jvp_jaxpr,
       name=f"{name}_jvp",
-      grid_mapping=jvp_grid_mapping,
+      in_shapes=(*in_shapes, *in_shapes),
+      out_shapes=(*out_shapes, *out_shapes),
+      grid_mapping=grid_mapping.replace(block_mappings=jvp_bms),
       interpret=interpret,
       debug=debug,
       input_output_aliases=(),
@@ -359,38 +364,40 @@ def _pallas_call_jvp_rule(primals, tangents, *, jaxpr, name,
   return out_primals, out_tangents
 ad.primitive_jvps[pallas_call_p] = _pallas_call_jvp_rule
 
-def _batch_block_mapping(grid_mapping: GridMapping,
-                         axis_size: int,
-                         aval: jax_core.ShapedArray,
+def _batch_block_mapping(grid_mapping: GridMapping, aval: jax_core.ShapedArray,
                          dim: int | batching.NotMapped,
-                         block_mapping: BlockMapping) -> BlockMapping:
+                         block_mapping: BlockMapping | None) -> BlockMapping:
   def _block_map_function(new_idx, *args):
-    indices = jax_core.eval_jaxpr(block_mapping.index_map_jaxpr.jaxpr,
-                                  block_mapping.index_map_jaxpr.consts,
-                                  *args)
+    if block_mapping is None:
+      indices = [0] * len(aval.shape)
+    else:
+      indices = jax_core.eval_jaxpr(block_mapping.index_map_jaxpr.jaxpr,
+                                    block_mapping.index_map_jaxpr.consts,
+                                    *args)
     if dim is not batching.not_mapped:
       indices.insert(dim, new_idx)
     return tuple(indices)
   i32_aval = jax_core.ShapedArray((), jnp.int32)
-  idx_avals = [i32_aval, *block_mapping.index_map_jaxpr.in_avals]
+  if block_mapping is None:
+    idx_avals = [i32_aval] * (len(grid_mapping.grid) + 1)
+  else:
+    idx_avals = [i32_aval, *block_mapping.index_map_jaxpr.in_avals]
   with grid_mapping.trace_env():
     block_mapping_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
         lu.wrap_init(_block_map_function), idx_avals)
-  shape = block_mapping.block_shape
+  shape = aval.shape if block_mapping is None else block_mapping.block_shape
   if dim is batching.not_mapped:
     new_block_shape = shape
-    new_array_shape_dtype = block_mapping.array_shape_dtype
   else:
     new_block_shape = tuple_insert(shape, dim, pallas_core.mapped)
-    new_array_shape_dtype = jax.ShapeDtypeStruct(
-        tuple_insert(block_mapping.array_shape_dtype.shape,
-                     dim,
-                     axis_size),
-        block_mapping.array_shape_dtype.dtype)
-
   jaxpr = jax_core.ClosedJaxpr(block_mapping_jaxpr, consts)
+  if block_mapping is None:
+    return BlockMapping(
+        block_shape=new_block_shape,
+        index_map_jaxpr=jaxpr,
+        indexing_mode=pallas_core.blocked,
+    )
   return block_mapping.replace(block_shape=new_block_shape,
-                               array_shape_dtype=new_array_shape_dtype,
                                index_map_jaxpr=jaxpr)
 
 
@@ -429,6 +436,8 @@ def _batch_with_explicit_loop(
     *,
     jaxpr: jax_core.Jaxpr,
     name: str,
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     grid_mapping: GridMapping,
     input_output_aliases: tuple[tuple[int, int], ...],
     debug: bool,
@@ -446,8 +455,7 @@ def _batch_with_explicit_loop(
   to the current iteration index and dynamic_updates an (initially empty) output
   allocation.
   """
-  # TODO(necula): cleanup
-  out_shapes = grid_mapping.out_shapes
+
   if not dims:
     raise NotImplementedError("vmapping pallas_call with no arguments.")
 
@@ -493,10 +501,13 @@ def _batch_with_explicit_loop(
                 axis=dim,
             )
         )
+
     batch_out = pallas_call_p.bind(
         *batch_args,
         jaxpr=jaxpr,
         name=name,
+        in_shapes=in_shapes,
+        out_shapes=out_shapes,
         grid_mapping=grid_mapping,
         input_output_aliases=input_output_aliases,
         debug=debug,
@@ -524,14 +535,15 @@ def _pallas_call_batching_rule(
     *,
     jaxpr: jax_core.Jaxpr,
     name: str,
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     grid_mapping: GridMapping,
     input_output_aliases: tuple[tuple[int, int], ...],
     debug: bool,
     interpret: bool,
     compiler_params: Any,
 ):
-  # TODO(necula): cleanup
-  out_shapes = grid_mapping.out_shapes
+
   def _maybe_squeeze_out_bdim(
       x: jax.Array, bdim: int | batching.NotMapped
   ) -> jax.Array:
@@ -548,6 +560,8 @@ def _pallas_call_batching_rule(
         *args,
         jaxpr=jaxpr,
         name=name,
+        in_shapes=in_shapes,
+        out_shapes=out_shapes,
         grid_mapping=grid_mapping,
         input_output_aliases=input_output_aliases,
         debug=debug,
@@ -579,6 +593,8 @@ def _pallas_call_batching_rule(
         dims=dynamic_grid_dims + dims,
         jaxpr=jaxpr,
         name=name,
+        in_shapes=in_shapes,
+        out_shapes=out_shapes,
         grid_mapping=grid_mapping,
         input_output_aliases=input_output_aliases,
         debug=debug,
@@ -611,6 +627,8 @@ def _pallas_call_batching_rule(
           dims=scalar_bdims + bdims,
           jaxpr=jaxpr,
           name=name,
+          in_shapes=in_shapes,
+          out_shapes=out_shapes,
           grid_mapping=grid_mapping,
           input_output_aliases=input_output_aliases,
           debug=debug,
@@ -636,6 +654,7 @@ def _pallas_call_batching_rule(
   all_dims = list(dims) + [0] * len(out_shapes)
 
   num_index_operands = grid_mapping.num_index_operands
+  num_constant_operands = grid_mapping.num_constant_operands
   num_scratch_operands = grid_mapping.num_scratch_operands
 
   # Only add a batch dimension for the avals that actually have a grid mapping.
@@ -643,27 +662,37 @@ def _pallas_call_batching_rule(
   # operands (the last in the list).
   avals_to_batch = avals[num_index_operands:(len(avals) - num_scratch_operands)]
   batched_block_mappings = map(
-      partial(_batch_block_mapping, grid_mapping, axis_size),
+      partial(_batch_block_mapping, grid_mapping),
       avals_to_batch,
       all_dims[num_index_operands:],
       block_mappings,
   )
 
-  index_map_tree_args, index_map_tree_kwargs = grid_mapping.index_map_tree.unflatten(grid_mapping.index_map_avals)
-  assert not index_map_tree_kwargs
-  batched_index_map_args = (jax_core.ShapedArray((), jnp.int32),) + index_map_tree_args
-  batched_index_map_avals, batched_index_map_tree = tree_util.tree_flatten((batched_index_map_args, {}))
+  # TODO(necula): should fix in_shapes to include the consts
+  dims_no_consts = (
+      dims[:num_index_operands] +
+      dims[num_index_operands + num_constant_operands:]
+  )
+  batched_in_shapes = tuple(
+      jax.ShapeDtypeStruct(x.shape if dim is batching.not_mapped else
+                           tuple_insert(x.shape, dim, axis_size),
+                           x.dtype)
+      for x, dim in zip(in_shapes, dims_no_consts))
+  batched_out_shapes = tuple(
+      jax.ShapeDtypeStruct(tuple_insert(x.shape, 0, axis_size), x.dtype)
+      for x in out_shapes)
+
   batched_grid_mapping = grid_mapping.replace(
       grid=(axis_size, *grid_mapping.grid),
       block_mappings=tuple(batched_block_mappings),
-      index_map_avals=batched_index_map_avals,
-      index_map_tree=batched_index_map_tree,
-      vmapped_dims=(0,) + tuple(a + 1 for a in grid_mapping.vmapped_dims))
+      mapped_dims=(0,) + tuple(a + 1 for a in grid_mapping.mapped_dims))
   out = pallas_call_p.bind(
       *dynamic_grid_args,
       *args,
       jaxpr=jaxpr,
       name=f"batched_{name}",
+      in_shapes=batched_in_shapes,
+      out_shapes=batched_out_shapes,
       grid_mapping=batched_grid_mapping,
       input_output_aliases=input_output_aliases,
       debug=debug,
@@ -698,6 +727,7 @@ def pallas_call_checkify_rule(error: checkify.Error,
                               interpret: bool,
                               input_output_aliases: tuple[tuple[int, int], ...],
                               grid_mapping: GridMapping,
+                              out_shapes,
                               **kwargs):
   # We implement the checkify rule in 4 steps:
   # 1) First, trace the kernel body to get the expected error shapes.
@@ -803,25 +833,37 @@ def pallas_call_checkify_rule(error: checkify.Error,
 
   # Prepare pallas_call inputs. We need to create new block specs
   # for the new error inputs and outputs.
-  error_block_specs = [pallas_core.BlockSpec(None, None)] * len(shaped_err_avals)
+  scalar_avals = map(checkify.get_shaped_aval, scalars)
+  error_block_specs = [pallas_core.BlockSpec(
+    index_map=lambda *args: (0,) * len(error.shape),
+    block_shape=error.shape)
+    for error in shaped_err_avals]
   error_paths, _ = unzip2(tree_util.tree_flatten_with_path(error_block_specs)[0])
+  grid_avals = [
+      jax_core.ShapedArray((), jnp.dtype("int32"))] * len(grid_mapping.grid)
+  scalar_ref_avals = [
+      pallas_core.AbstractMemoryRef(
+          jax_core.ShapedArray(aval.shape, aval.dtype),
+          pallas_core.MemorySpace.INDEX)
+      for aval in scalar_avals]
+  grid_tree = tree_util.tree_structure(((*grid_avals, *scalar_avals), {}))
   error_block_mappings = map(
         partial(
             pallas_core._convert_block_spec_to_block_mapping,
-            index_map_avals=grid_mapping.index_map_avals,
-            index_map_tree=grid_mapping.index_map_tree,
+            (*grid_avals, *scalar_ref_avals),
+            in_tree=grid_tree,
             grid=grid_mapping.grid,
-            mapped_dims=grid_mapping.vmapped_dims,
+            mapped_dims=grid_mapping.mapped_dims,
             what="error"),
-        error_block_specs, error_paths, shaped_err_avals)
+        error_block_specs, error_paths, error_memref_aval)
   input_block_mappings, output_block_mappings = split_list(
       grid_mapping.block_mappings, [num_kernel_inputs,])
   grid_mapping_with_error = grid_mapping.replace(
       block_mappings=(*error_block_mappings, *input_block_mappings,
-                      *error_block_mappings, *output_block_mappings),
-      num_inputs=grid_mapping.num_inputs + len(error_block_mappings),
-      num_outputs=grid_mapping.num_outputs + len(error_block_mappings)
+                      *error_block_mappings, *output_block_mappings)
   )
+  error_out_shapes = tuple(
+      jax.ShapeDtypeStruct(e.shape, e.dtype) for e in shaped_err_avals)
   # Bump all input_output_aliases by num_err_vals to make room for error
   # TODO(justinfu): Don't bump scalars here.
   input_output_aliases = tuple(
@@ -830,11 +872,17 @@ def pallas_call_checkify_rule(error: checkify.Error,
       (i+num_scalars, i) for i in range(num_err_vals)) + input_output_aliases
 
   new_vals_in = [*scalars, *err_vals, *args]
+  new_input_shapes = tuple(
+      jax.ShapeDtypeStruct(x.shape, x.dtype) for x in [
+          *scalars, *shaped_err_avals, *args])
+  del kwargs['in_shapes']
   result = pallas_call_p.bind(*dynamic_grid_bounds, *new_vals_in,
     jaxpr=final_jaxpr,
     interpret=interpret,
     grid_mapping=grid_mapping_with_error,
     input_output_aliases=input_output_aliases_with_error,
+    in_shapes=new_input_shapes,
+    out_shapes=error_out_shapes + out_shapes,
     **kwargs)
   errors, results = split_list(result, [num_err_vals])
   # TODO(b/350593266): Remove line below once we support ()-shaped scalars.
@@ -844,24 +892,24 @@ def pallas_call_checkify_rule(error: checkify.Error,
 checkify.error_checks[pallas_call_p] = pallas_call_checkify_rule
 
 @weakref_lru_cache
-def _trace_kernel_to_jaxpr(fun: Callable, grid_spec: GridSpec,
-                           flat_in_avals: Sequence[jax_core.AbstractValue],
-                           flat_out_avals: Sequence[jax_core.AbstractValue],
-                           in_tree: tree_util.PyTreeDef,
-                           in_paths: Sequence[tree_util.KeyPath],
-                           out_tree: tree_util.PyTreeDef,
-                           out_paths: Sequence[tree_util.KeyPath],
-                           interpret: bool):
+def _trace_to_jaxpr(fun: Callable, grid_spec: GridSpec,
+                    flat_in_avals: Sequence[jax_core.AbstractValue],
+                    flat_out_avals: Sequence[jax_core.AbstractValue],
+                    in_tree: tree_util.PyTreeDef,
+                    in_paths: Sequence[tree_util.KeyPath],
+                    out_tree: tree_util.PyTreeDef,
+                    out_paths: Sequence[tree_util.KeyPath],
+                    interpret: bool):
   avals, grid_mapping = grid_spec.get_grid_mapping(flat_in_avals, in_tree, in_paths,
                                                    flat_out_avals, out_tree, out_paths)
   if interpret:
     avals = jax.tree_util.tree_map(_logical_aval_to_interpret_mode_aval, avals)
   jaxpr_flat_avals, jaxpr_in_tree = tree_util.tree_flatten(avals)
-  wrapped_kernel_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
+  wrapped_fun, out_tree_thunk = api_util.flatten_fun_nokwargs(
       lu.wrap_init(fun), jaxpr_in_tree)
   debug = pe.debug_info(fun, jaxpr_in_tree, out_tree_thunk, False, "pallas_call")
   with grid_mapping.trace_env():
-    jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_kernel_fun,
+    jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_fun,
                                                      jaxpr_flat_avals, debug)
     if consts:
       # Pad ``block_mappings`` to account for the hoisted constants.
@@ -872,14 +920,20 @@ def _trace_kernel_to_jaxpr(fun: Callable, grid_spec: GridSpec,
           index=grid_mapping.num_index_operands,
           make_abstract_ref=lambda aval: pallas_core.AbstractMemoryRef(aval, None))
       num_constant_operands = len(consts)
+      # TODO(necula): refactor grid_mapping to remove this code duplication
+      grid_avals = [jax_core.ShapedArray((), jnp.dtype("int32"))] * len(grid_mapping.grid)
+      if grid_mapping.num_index_operands:
+        grid_avals += flat_in_avals[:grid_mapping.num_index_operands]  # type: ignore
+      # Create args, kwargs pytree def
+      grid_tree = tree_util.tree_structure((tuple(grid_avals), {}))
       const_block_mappings = []
       for c_idx, c in enumerate(consts):
         const_block_mapping = pallas_core._convert_block_spec_to_block_mapping(
+            grid_avals,
             pallas_core.BlockSpec(None, None),
             path=(tree_util.SequenceKey(c_idx),),
-            array_aval=jax_core.ShapedArray(c.shape, c.dtype),
-            index_map_avals=grid_mapping.index_map_avals,
-            index_map_tree=grid_mapping.index_map_tree,
+            aval=jax_core.ShapedArray(c.shape, c.dtype),
+            in_tree=grid_tree,
             grid=grid_mapping.grid,
             mapped_dims=(),
             what="consts",
@@ -890,12 +944,7 @@ def _trace_kernel_to_jaxpr(fun: Callable, grid_spec: GridSpec,
           block_mappings=(*const_block_mappings, *grid_mapping.block_mappings),
           num_constant_operands=num_constant_operands,
       )
-  kernel_out_tree = out_tree_thunk()
-  if kernel_out_tree != tree_util.tree_structure(None):
-    raise ValueError(
-        "The kernel function in a pallas_call should return None. "
-        f"Found a PyTree: {kernel_out_tree}")
-  return grid_mapping, jaxpr, consts
+  return grid_mapping, jaxpr, consts, out_tree_thunk()
 
 def _extract_function_name(f: Callable, name: str | None) -> str:
   if name is None:
@@ -997,7 +1046,7 @@ def pallas_call(
     *,
     grid_spec: GridSpec | None = None,
     debug: bool = False,
-    grid: Grid = (),
+    grid: Grid | None = None,
     in_specs: BlockSpecTree = no_block_spec,
     out_specs: BlockSpecTree = no_block_spec,
     input_output_aliases: dict[int, int] = {},
@@ -1019,7 +1068,8 @@ def pallas_call(
     debug: if True, Pallas prints various intermediate forms of the kernel
       as it is being processed.
     grid: the iteration space, as a tuple of integers. The kernel is executed
-      as many times as ``prod(grid)``.
+      as many times as ``prod(grid)``. The default value ``None`` is equivalent
+      to ``()``.
       See details at :ref:`pallas_grid`.
     in_specs: a PyTree of :class:`jax.experimental.pallas.BlockSpec` with
       a structure matching that of the positional arguments.
@@ -1049,14 +1099,10 @@ def pallas_call(
   name = _extract_function_name(f, name)
   if compiler_params is None:
     compiler_params = {}
-  if grid_spec is not None and not (
-      grid == () and
-      in_specs is no_block_spec and
-      out_specs is no_block_spec):
-    raise ValueError("Cannot specify both grid_spec and at the same time any of grid, in_specs, out_specs.")
+  if grid is not None and grid_spec is not None:
+    raise ValueError("Cannot specify both grid and grid_spec at the same time.")
   if grid_spec is None:
     grid_spec = GridSpec(grid, in_specs, out_specs)
-  del grid, in_specs, out_specs
   grid_spec, dynamic_grid_bounds = grid_spec.unzip_dynamic_grid_bounds()
   # TODO(necula): this canonicalization may be convenient for some usage
   # but it is lossy, because it prevents expressing functions that return
@@ -1075,10 +1121,14 @@ def pallas_call(
                           for a in flat_args)
     flat_out_avals = tuple(jax_core.ShapedArray(v.shape, v.dtype)
                            for v in flat_out_shapes)
-    # TODO(necula): check that input_output_aliases is well-formed: no duplicates, etc.
-    grid_mapping, jaxpr, consts = _trace_kernel_to_jaxpr(
+    # TODO(necula): check that input_output_aliases is well-formed: shapes match, no duplicates, etc.
+    grid_mapping, jaxpr, consts, f_out_tree = _trace_to_jaxpr(
         f, grid_spec, flat_in_avals, flat_out_avals, in_tree, in_paths,
         out_tree, out_paths, interpret=interpret)
+    if f_out_tree != tree_util.tree_flatten(None)[1]:
+      raise ValueError(
+          "The kernel function in a pallas_call should return None. "
+          f"Found a PyTree: {f_out_tree}")
     for i_idx, o_idx in input_output_aliases.items():
       if i_idx not in range(len(flat_in_avals)):
         raise ValueError(
@@ -1104,7 +1154,9 @@ def pallas_call(
     out_flat = pallas_call_p.bind(
         *dynamic_grid_bounds, *index_args, *consts, *rest_args,
         jaxpr=jaxpr, name=name,
-        debug=debug,
+        in_shapes=tuple(jax.ShapeDtypeStruct(a.shape, a.dtype)
+                        for a in flat_args),
+        out_shapes=tuple(flat_out_shapes), debug=debug,
         interpret=interpret,
         grid_mapping=grid_mapping,
         input_output_aliases=tuple(input_output_aliases.items()),
diff --git a/jax/_src/pallas/triton/lowering.py b/jax/_src/pallas/triton/lowering.py
index 65300ec5f..57fcc2808 100644
--- a/jax/_src/pallas/triton/lowering.py
+++ b/jax/_src/pallas/triton/lowering.py
@@ -112,8 +112,10 @@ class LoweringError(Exception):
 
 
 def _eval_index_map(
-    ctx: ModuleContext, idx, block_mapping: BlockMapping
+    ctx: ModuleContext, idx, block_mapping: BlockMapping | None
 ):
+  if block_mapping is None:
+    return None
   block_indices = lower_jaxpr_to_triton_ir(
       ctx, block_mapping.index_map_jaxpr.jaxpr, None, *idx
   )
@@ -185,13 +187,13 @@ def _process_grid_to_3d_grid(grid_mapping: GridMapping):
 
   # Preserve grid order provided to pallas_call
   for i, s in enumerate(grid_mapping.grid):
-    if i not in grid_mapping.vmapped_dims:
+    if i not in grid_mapping.mapped_dims:
       launch_grid.append(s)
       launch_grid_to_pallas_grid.append(i)
 
   # For mapped dims, iterate from inner to outer. This follows the pallas_call
   # batching rule that prepends the vmapped dimension.
-  for dim in reversed(grid_mapping.vmapped_dims):
+  for dim in reversed(grid_mapping.mapped_dims):
     s = grid_mapping.grid[dim]
     launch_grid.append(s)
     launch_grid_to_pallas_grid.append(dim)
@@ -285,7 +287,7 @@ def lower_jaxpr_to_triton_module(
       local_program_ids = [
           pid
           for i, pid in enumerate(program_ids)
-          if i not in grid_mapping.vmapped_dims
+          if i not in grid_mapping.mapped_dims
       ]
       ctx = ModuleContext(
           name, grid_mapping, local_program_ids, mlir.TracebackCaches(), platform
@@ -295,7 +297,7 @@ def lower_jaxpr_to_triton_module(
             "Scalar prefetch not supported in Triton lowering."
         )
       for bm in grid_mapping.block_mappings:
-        if not isinstance(bm.indexing_mode, Blocked):
+        if bm is not None and not isinstance(bm.indexing_mode, Blocked):
           raise NotImplementedError(
               "Only Blocked indexing mode is supported in Triton lowering."
           )
@@ -303,14 +305,20 @@ def lower_jaxpr_to_triton_module(
           functools.partial(_eval_index_map, ctx, program_ids),
           grid_mapping.block_mappings,
       )
+      consts_shapes = [
+          jax.ShapeDtypeStruct(v.aval.shape, v.aval.dtype)
+          for v in jaxpr.invars[grid_mapping.num_index_operands:grid_mapping.num_index_operands + grid_mapping.num_constant_operands]
+      ]
       block_infos = [
           BlockInfo(
               jax.ShapeDtypeStruct(shape_dtype.shape, shape_dtype.dtype),
               start_idx,
               block_mapping.block_shape,
           )
+          if block_mapping is not None
+          else None
           for shape_dtype, block_mapping, start_idx in zip(
-              in_out_shapes,
+              (*consts_shapes, *in_out_shapes),
               grid_mapping.block_mappings,
               start_indices,
           )
diff --git a/jax/_src/pallas/triton/pallas_call_registration.py b/jax/_src/pallas/triton/pallas_call_registration.py
index 052e27aae..800b328f2 100644
--- a/jax/_src/pallas/triton/pallas_call_registration.py
+++ b/jax/_src/pallas/triton/pallas_call_registration.py
@@ -19,10 +19,12 @@ from __future__ import annotations
 import io
 from typing import Any
 
+import jax
 from jax import core as jax_core
 from jax._src.interpreters import mlir
 from jax._src.lib.mlir import ir
 from jax._src.pallas import core as pallas_core
+from jax._src.pallas.pallas_call import pallas_call_p
 from jax._src.pallas.triton import lowering
 
 
@@ -43,16 +45,30 @@ def pallas_call_lowering(
     *in_nodes,
     jaxpr: jax_core.Jaxpr,
     name: str,
+    in_shapes: tuple[jax.ShapeDtypeStruct, ...],
+    out_shapes: tuple[jax.ShapeDtypeStruct, ...],
     interpret: bool,
     debug: bool,
     input_output_aliases: tuple[tuple[int, int], ...],
     grid_mapping: pallas_core.GridMapping,
     compiler_params: dict[str, Any],
 ):
-  del interpret
-  # TODO(necula): cleanup
-  in_shapes = grid_mapping.in_shapes
-  out_shapes = grid_mapping.out_shapes
+  if interpret:
+    # TODO(necula): is this branch still needed?
+    return mlir.lower_fun(pallas_call_p.impl, multiple_results=True)(
+        ctx,
+        *in_nodes,
+        jaxpr=jaxpr,
+        name=name,
+        out_shapes=out_shapes,
+        in_shapes=in_shapes,
+        interpret=interpret,
+        debug=debug,
+        input_output_aliases=input_output_aliases,
+        grid_mapping=grid_mapping,
+        compiler_params=compiler_params,
+    )
+
   if grid_mapping.num_dynamic_grid_bounds:
     raise NotImplementedError(
         "dynamic grid bounds not supported in the Triton backend"
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index 4996621f7..e6918f459 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -19,7 +19,6 @@ import ctypes
 import dataclasses
 import functools
 import itertools
-import math
 import os
 import pathlib
 import subprocess
@@ -43,6 +42,7 @@ from jaxlib.mlir.dialects import nvvm
 from jaxlib.mlir.passmanager import PassManager
 import numpy as np
 
+from . import dsl as mgpu
 from . import profiler
 from . import utils
 
@@ -63,7 +63,7 @@ TMA_DESCRIPTOR_BYTES = 128
 TMA_DESCRIPTOR_ALIGNMENT = 64
 
 
-c = utils.c  # This is too common to fully qualify.
+c = mgpu.c  # This is too common to fully qualify.
 
 
 RUNTIME_PATH = None
@@ -143,13 +143,13 @@ class TileTransform(MemRefTransform):
     tiling_rank = len(self.tiling)
     tiled_rank = untiled_rank + tiling_rank
     for t, d in zip(self.tiling[::-1], range(untiled_rank)[::-1]):
-      ref = utils.memref_unfold(ref, d, (None, t))
+      ref = mgpu.memref_unfold(ref, d, (None, t))
     permutation = (
         *range(untiled_rank - tiling_rank),
         *range(untiled_rank - tiling_rank, tiled_rank, 2),
         *range(untiled_rank - tiling_rank + 1, tiled_rank, 2),
     )
-    return utils.memref_transpose(ref, permutation)
+    return mgpu.memref_transpose(ref, permutation)
 
   def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]:
     index = ir.IndexType.get()
@@ -193,7 +193,7 @@ class TransposeTransform(MemRefTransform):
       raise ValueError("Permutation must be a permutation")
 
   def apply(self, ref: ir.Value) -> ir.Value:
-    return utils.memref_transpose(ref, self.permutation)
+    return mgpu.memref_transpose(ref, self.permutation)
 
   def transform_index(self, idx: Sequence[ir.Value]) -> tuple[ir.Value, ...]:
     return tuple(idx[p] for p in self.permutation)
@@ -209,7 +209,7 @@ OnDeviceProfiler = profiler.OnDeviceProfiler
 class LaunchContext:
   launch_op: gpu.LaunchOp
   gmem_scratch_ptr: ir.Value
-  cluster_size: tuple[int, int, int]
+  cluster_size: tuple[int, ...]
   profiler: OnDeviceProfiler | None = None
   next_scratch_offset: int = 0
   host_scratch_init: list[Callable[[ir.Value], None]] = dataclasses.field(
@@ -261,6 +261,33 @@ class LaunchContext:
     gep.move_after(self.gmem_scratch_ptr.owner)
     return device_init(gep.result)
 
+  def cluster_collective_mask(self, collective: gpu.Dimension):
+    # We first compute the linearized index of the slice along the collective
+    # dim that contains the current block. Then, the mask is a sequence of 1s
+    # strided by the position of the collective dim, shifted left by the linear
+    # slice index.
+    # TODO(apaszke): Make sure this gets hoisted outside of any loops.
+    # If not, we might need to do it manually.
+    i32 = ir.IntegerType.get_signless(32)
+    stride = 1
+    mask_shift = c(0, i32)
+    collective_stride = None
+    for cluster_dim in gpu.Dimension:
+      if self.cluster_size[cluster_dim] == 1:
+        continue
+      if cluster_dim != collective:
+        dim_idx = arith.index_castui(i32, gpu.cluster_block_id(cluster_dim))
+        mask_shift = arith.addi(
+            mask_shift, arith.muli(dim_idx, c(stride, i32)),
+        )
+      else:
+        collective_stride = stride
+      stride *= self.cluster_size[cluster_dim]
+    mask_unshifted = 0
+    for i in range(self.cluster_size[collective]):
+      mask_unshifted |= 1 << (i * collective_stride)
+    return arith.shli(c(mask_unshifted, i32), mask_shift)
+
   def _get_tma_desc(
       self,
       gmem_ref,
@@ -319,7 +346,7 @@ class LaunchContext:
       dst_ref,
       gmem_slice: Any = (),
       gmem_transform: MemRefTransform | tuple[MemRefTransform, ...] = (),
-      barrier: utils.BarrierRef | None = None,
+      barrier: mgpu.Barrier | None = None,
       swizzle: int | None = None,
       arrive: bool | None = None,
       uniform: bool = True,
@@ -332,7 +359,7 @@ class LaunchContext:
     src_ref_ty = ir.MemRefType(src_ref.type)
     dst_ref_ty = ir.MemRefType(dst_ref.type)
     element_type = src_ref_ty.element_type
-    element_bytewidth = utils.bytewidth(element_type)
+    element_bytewidth = mgpu.bytewidth(element_type)
     if element_type != dst_ref_ty.element_type:
       raise ValueError(
           f"Expected same element type, got {element_type} and"
@@ -375,7 +402,7 @@ class LaunchContext:
       slice_shape = t.transform_shape(slice_shape)
     for dim, squeezed in enumerate(is_squeezed):
       if squeezed:
-        smem_ref = utils.memref_unsqueeze(smem_ref, dim)
+        smem_ref = mgpu.memref_unsqueeze(smem_ref, dim)
     smem_ref_ty = ir.MemRefType(smem_ref.type)
 
     if slice_shape != tuple(smem_ref_ty.shape):
@@ -407,13 +434,13 @@ class LaunchContext:
       dyn_base_indices[collective_slice_dim] = arith.addi(
           dyn_base_indices[collective_slice_dim], block_offset,
       )
-      smem_ref = utils.memref_slice(
+      smem_ref = mgpu.memref_slice(
           smem_ref,
           (slice(None),) * collective_slice_dim
-          + (utils.ds(block_offset, slice_shape[collective_slice_dim]),),
+          + (mgpu.ds(block_offset, slice_shape[collective_slice_dim]),),
       )
       multicast_mask = arith.trunci(
-          i16, utils.cluster_collective_mask(self.cluster_size, collective)
+          i16, self.cluster_collective_mask(collective)
       )
     else:
       multicast_mask = None
@@ -428,7 +455,7 @@ class LaunchContext:
     ]
 
     uniform_ctx = (
-        functools.partial(utils.single_thread, per_block=False)
+        functools.partial(mgpu.single_thread, per_block=False)
         if uniform
         else contextlib.nullcontext
     )
@@ -479,121 +506,36 @@ T = TypeVar('T')
 class Union(Generic[T]):
   members: Sequence[T]
 
-  def __iter__(self):
-    return iter(self.members)
-
-@dataclasses.dataclass(frozen=True)
-class TMABarrier:
-  num_barriers: int = 1
-
-@dataclasses.dataclass(frozen=True)
-class Barrier:
-  arrival_count: int
-  num_barriers: int = 1
-
-@dataclasses.dataclass(frozen=True)
-class ClusterBarrier:
-  collective_dims: Sequence[gpu.Dimension]
-  num_barriers: int = 1
-
 
 def _count_buffer_bytes(shape_dtype: jax.ShapeDtypeStruct) -> int:
   return np.prod(shape_dtype.shape) * np.dtype(shape_dtype.dtype).itemsize
 
 
 def _construct_smem_reftree(
-    cluster_shape: tuple[int, int, int],
-    dynamic_smem: ir.Value,
-    smem_buffers: ShapeTree,
-    dynamic_smem_offset: int = 0,
-) -> RefTree:
+    dynamic_smem: ir.Value, smem_buffers: ShapeTree) -> RefTree:
   index = ir.IndexType.get()
-  i8 = ir.IntegerType.get_signless(8)
-  ptr = ir.Type.parse("!llvm.ptr")
   smem = ir.Attribute.parse("#gpu.address_space<workgroup>")
-  flat_ref_tys, smem_buffer_tree = jax.tree.flatten(
-      smem_buffers, is_leaf=lambda x: isinstance(x, Union)
-  )
+  flat_ref_tys, smem_buffer_tree = jax.tree.flatten(smem_buffers)
   smem_refs = []
+  dynamic_smem_offset = 0
   for ref_ty in flat_ref_tys:
-    def get_barrier_ptr(num_barriers: int) -> ir.Value:
-      nonlocal dynamic_smem_offset
-      smem_base_ptr = utils.memref_ptr(dynamic_smem, memory_space=3)
-      barrier_base_ptr = llvm.getelementptr(
-          ptr, smem_base_ptr, [], [dynamic_smem_offset], i8
-      )
-      dynamic_smem_offset += num_barriers * MBARRIER_BYTES
-      return barrier_base_ptr
-    match ref_ty:
-      case Union(members):
-        member_trees = [
-            _construct_smem_reftree(cluster_shape, dynamic_smem, m, dynamic_smem_offset)
-            for m in members
-        ]
-        # TODO(apaszke): This is quadratic, but it shouldn't matter for now...
-        dynamic_smem_offset += _smem_tree_size(ref_ty)
-        ref = Union(member_trees)
-      case TMABarrier(num_barriers):
-        ref = utils.BarrierRef.initialize(
-            get_barrier_ptr(num_barriers), num_barriers, arrival_count=1
-        )
-      case Barrier(arrival_count, num_barriers):
-        ref = utils.BarrierRef.initialize(
-            get_barrier_ptr(num_barriers),
-            num_barriers,
-            arrival_count=arrival_count,
-        )
-      case ClusterBarrier(collective_dims, num_barriers):
-        ref = utils.CollectiveBarrierRef.initialize(
-            get_barrier_ptr(num_barriers),
-            num_barriers,
-            collective_dims,
-            cluster_shape,
-        )
-      case _:
-        mlir_dtype = mlir.dtype_to_ir_type(ref_ty.dtype)
-        tile_smem = memref.view(
-            ir.MemRefType.get(ref_ty.shape, mlir_dtype, memory_space=smem),
-            dynamic_smem, c(dynamic_smem_offset, index), [],
-        )
-        dynamic_smem_offset += _count_buffer_bytes(ref_ty)
-        ref = tile_smem
-    smem_refs.append(ref)
+    mlir_dtype = mlir.dtype_to_ir_type(ref_ty.dtype)
+    tile_smem = memref.view(
+        ir.MemRefType.get(ref_ty.shape, mlir_dtype, memory_space=smem),
+        dynamic_smem, c(dynamic_smem_offset, index), [],
+    )
+    dynamic_smem_offset += _count_buffer_bytes(ref_ty)
+    smem_refs.append(tile_smem)
   return jax.tree.unflatten(smem_buffer_tree, smem_refs)
 
 
-MBARRIER_BYTES = 8
-
-
-def _smem_tree_size(smem_buffers: ShapeTree) -> int:
-  leaves = jax.tree.leaves(
-      smem_buffers, is_leaf=lambda x: isinstance(x, Union)
-  )
-  size = 0
-  for l in leaves:
-    match l:
-      case Union(members):
-        size += max(_smem_tree_size(s) for s in members)
-      case (
-          TMABarrier(num_barriers)
-          | ClusterBarrier(_, num_barriers=num_barriers)
-          | Barrier(_, num_barriers=num_barriers)
-      ):
-        if size % MBARRIER_BYTES:
-          raise NotImplementedError("Misaligned barrier allocation")
-        size += num_barriers * MBARRIER_BYTES
-      case _:
-        size += _count_buffer_bytes(l)
-  return size
-
-
 # TODO(apaszke): Inline this
 @contextlib.contextmanager
 def _launch(
     token,
-    grid: tuple[int, int, int],
-    cluster: tuple[int, int, int],
-    block: tuple[int, int, int],
+    grid,
+    cluster,
+    block,
     scratch_arr,
     smem_buffers: ShapeTree | Union[ShapeTree],
     profiler_spec: profiler.ProfilerSpec | None = None,
@@ -607,16 +549,24 @@ def _launch(
   grid_vals = [c(i, index) for i in grid]
   block_vals = [c(i, index) for i in block]
 
-  user_smem_bytes = _smem_tree_size(smem_buffers)
+  if isinstance(smem_buffers, Union):
+    smem_disjoint_live_buffers_collections = smem_buffers.members
+    compute_smem_bytes = max(
+        sum(_count_buffer_bytes(l) for l in jax.tree.leaves(s))
+            for s in smem_buffers.members)
+  else:
+    smem_disjoint_live_buffers_collections = [smem_buffers]
+    compute_smem_bytes = sum(
+        _count_buffer_bytes(l) for l in jax.tree.leaves(smem_buffers))
 
-  smem_bytes = user_smem_bytes
+  smem_bytes = compute_smem_bytes
   if profiler_spec is not None:
     smem_bytes += profiler_spec.smem_bytes(block=block)
 
   # TODO(cperivol): Query the shared memory size programmatically.
   if smem_bytes > 228 * 1024:
     raise ValueError(f"Mosaic GPU kernel exceeds available shared memory {smem_bytes=} > 228000")
-  if math.prod(cluster) != 1:
+  if cluster:
     if len(cluster) != 3:
       raise ValueError("Clusters must be 3D")
     cluster_kwargs = {
@@ -642,15 +592,11 @@ 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()
+    smem_ref_trees = []
+    for smem_live_buffers_collection in smem_disjoint_live_buffers_collections:
+      smem_ref_tree = _construct_smem_reftree(
+          dynamic_smem, smem_live_buffers_collection)
+      smem_ref_trees.append(smem_ref_tree)
 
     if profiler_spec:
       prof_smem = memref.view(
@@ -658,7 +604,7 @@ def _launch(
               (profiler_spec.smem_i32_elements(block=block),),
               i32, memory_space=smem,
           ),
-          dynamic_smem, c(user_smem_bytes, index), [],
+          dynamic_smem, c(compute_smem_bytes, index), [],
       )
       prof = profiler.OnDeviceProfiler(
           profiler_spec, prof_smem, maybe_prof_buffer
@@ -666,6 +612,11 @@ def _launch(
     else:
       prof = None
 
+    if isinstance(smem_buffers, Union):
+      smem_ref_tree: Union[RefTree] = Union(smem_ref_trees)
+    else:
+      smem_ref_tree: RefTree = smem_ref_trees[0] if smem_ref_trees else []
+
     ptr_ty = ir.Type.parse("!llvm.ptr")
     scratch_ptr = builtin.unrealized_conversion_cast([ptr_ty], [scratch_arr])
     yield LaunchContext(launch_op, scratch_ptr, cluster, prof), smem_ref_tree
@@ -676,9 +627,9 @@ def _launch(
 
 def _lower_as_gpu_kernel(
     body,
-    grid: tuple[int, int, int],
-    cluster: tuple[int, int, int],
-    block: tuple[int, int, int],
+    grid: tuple[int, ...],
+    cluster: tuple[int, ...],
+    block: tuple[int, ...],
     in_shapes: tuple[Any, ...],
     out_shape,
     smem_scratch_shape: ShapeTree | Union[ShapeTree],
@@ -686,6 +637,7 @@ def _lower_as_gpu_kernel(
 ):
   ptr_ty = ir.Type.parse("!llvm.ptr")
   token_ty = ir.Type.parse("!gpu.async.token")
+  i8 = ir.IntegerType.get_signless(8)
   i32 = ir.IntegerType.get_signless(32)
   i64 = ir.IntegerType.get_signless(64)
 
@@ -755,13 +707,13 @@ def _lower_as_gpu_kernel(
 
 def as_gpu_kernel(
     body,
-    grid: tuple[int, int, int],
-    block: tuple[int, int, int],
+    grid: tuple[int, ...],
+    block: tuple[int, ...],
     in_shape,
     out_shape,
     smem_scratch_shape: ShapeTree | Union[ShapeTree],
     prof_spec: profiler.ProfilerSpec | None = None,
-    cluster: tuple[int, int, int] = (1, 1, 1),
+    cluster: tuple[int, ...] = (),
 ):
   if isinstance(in_shape, list):
     in_shape = tuple(in_shape)
diff --git a/jax/experimental/mosaic/gpu/dsl.py b/jax/experimental/mosaic/gpu/dsl.py
index caf03a95c..bd8960c74 100644
--- a/jax/experimental/mosaic/gpu/dsl.py
+++ b/jax/experimental/mosaic/gpu/dsl.py
@@ -13,12 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from . import (
-    Barrier,
-    ClusterBarrier,
-    TMABarrier,
-    Union,
-)
 from .fragmented_array import (
     FragmentedArray,
     FragmentedLayout,
@@ -27,6 +21,8 @@ from .fragmented_array import (
     WGStridedFragLayout,
 )
 from .utils import (
+    Barrier,
+    BarrierArray,
     DynamicSlice,
     Partition,
     Partition1D,
diff --git a/jax/experimental/mosaic/gpu/examples/flash_attention.py b/jax/experimental/mosaic/gpu/examples/flash_attention.py
index 067584422..e680faffe 100644
--- a/jax/experimental/mosaic/gpu/examples/flash_attention.py
+++ b/jax/experimental/mosaic/gpu/examples/flash_attention.py
@@ -132,7 +132,7 @@ def build_kernel(
   v_scratch = jax.ShapeDtypeStruct(
       tile_shape((blocks.stages, blocks.kv, head_dim), tiling), jnp.float16
   )
-  smem_buffers_shape = [
+  smem_scratch_shape = [
       qo_scratch,
       k_scratch,
       v_scratch,
@@ -149,18 +149,19 @@ def build_kernel(
       k_gmem,
       v_gmem,
       out_gmem,
-      smem,
+      smem_scratch,
   ):
-    wg_idx = warpgroup_idx(sync=True)
-    smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = smem
-    qo_smem, k_smem, v_smem = smem_buffers
-    k_barriers, v_barriers, q_barriers = buffer_barriers
-    k_consumed_barrier, v_consumed_barrier = consumed_barriers
+    k_barriers = BarrierArray(blocks.stages)
+    v_barriers = BarrierArray(blocks.stages)
+    q_barriers = BarrierArray(compute_wgs_per_block)
+    k_consumed_barrier, v_consumed_barrier = BarrierArray(2, arrival_count=256)
+    schedule_barrier = BarrierArray(1, arrival_count=256)[0]
     @ctx.named_region("Schedule barrier")
     def perform_schedule_barrier():
       schedule_barrier.arrive()
       schedule_barrier.wait()
-
+    wg_idx = warpgroup_idx(sync=True)
+    qo_smem, k_smem, v_smem = smem_scratch
     qo_smem = memref_slice(qo_smem, arith.index_cast(index, wg_idx))
 
     @contextlib.contextmanager
@@ -360,12 +361,13 @@ def build_kernel(
       out_gmem,
       smem_scratch,
   ):
-    wg_idx = warpgroup_idx(sync=True)
-    (qo_smem, k_smem, v_smem), barriers, schedule_barrier = smem_scratch
+    barriers = BarrierArray(blocks.stages + wgs_per_block)
+    schedule_barrier = BarrierArray(1, arrival_count=256)[0]
     def perform_schedule_barrier():
       schedule_barrier.arrive()
       schedule_barrier.wait()
-
+    wg_idx = warpgroup_idx(sync=True)
+    qo_smem, k_smem, v_smem = smem_scratch
     qo_smem = memref_slice(qo_smem, arith.index_cast(index, wg_idx))
 
     @contextlib.contextmanager
@@ -401,8 +403,8 @@ def build_kernel(
 
     def kv_copy_init(slot, kv_seq_base):
       with single_thread(per_block=False):
-        txcount = 2 * blocks.kv * head_dim * bytewidth(f16)
-        barriers[slot].arrive_expect_tx(txcount)
+        txcount = c(2 * blocks.kv * head_dim * bytewidth(f16))
+        nvgpu.mbarrier_arrive_expect_tx(barriers.value, txcount, slot)
         k_tr = (
             mosaic_gpu.TileTransform(tiling),
             mosaic_gpu.TransposeTransform((0, 2, 1, 3, 4)),
@@ -537,23 +539,8 @@ def build_kernel(
   match impl:
     case Implementation.TWO_COMPUTE_WG:
       kernel = compute_only_kernel
-      smem_scratch_shape = (
-          smem_buffers_shape,
-          TMABarrier(blocks.stages + compute_wgs_per_block),
-          Barrier(arrival_count=256, num_barriers=1),
-      )
     case Implementation.TWO_COMPUTE_ONE_TMA_WG:
       kernel = tma_wg_kernel
-      smem_scratch_shape = (
-          smem_buffers_shape,
-          (
-              TMABarrier(blocks.stages),
-              TMABarrier(blocks.stages),
-              TMABarrier(compute_wgs_per_block),
-          ),
-          Barrier(arrival_count=256, num_barriers=2),
-          Barrier(arrival_count=256, num_barriers=1),
-      )
   return mosaic_gpu.as_gpu_kernel(
       kernel, grid, block, in_shape, out_shape, smem_scratch_shape, prof_spec
   )
diff --git a/jax/experimental/mosaic/gpu/examples/matmul.py b/jax/experimental/mosaic/gpu/examples/matmul.py
index ce24ed21f..53818f2ef 100644
--- a/jax/experimental/mosaic/gpu/examples/matmul.py
+++ b/jax/experimental/mosaic/gpu/examples/matmul.py
@@ -296,25 +296,27 @@ def build_kernel(
   smem_shape = mosaic_gpu.Union(
       [compute_scratch_shapes, epilogue_scratch_shapes])
 
-  def _main(ctx, a_device, b_device, c_device, smem):
-    (compute_smem, epilogue_smem), barriers = smem
+  def _main(ctx, a_device, b_device, c_device,
+            smem_union: mosaic_gpu.Union[mosaic_gpu.RefTree]):
+    compute_smem, epilogue_smem = smem_union.members
 
     memref.assume_alignment(c_device, 16)
 
+    barrier_group = BarrierArray(stages)
     m_start = arith.muli(c(block_tiling.m), gpu.block_id(gpu.Dimension.x))
     n_start = arith.muli(c(block_tiling.n), gpu.block_id(gpu.Dimension.y))
 
     def fetch(slot, ki):
-      barrier = barriers[slot]
+      barrier = barrier_group[slot]
       k_start = arith.muli(c(block_tiling.k), ki)
       lhs_tma_tile_bytes = int(np.prod(block_tiling.mk) * lhs_elem_bytes)
       rhs_tma_tile_bytes = int(np.prod(block_tiling.kn) * rhs_elem_bytes)
-      txcount = lhs_tma_tile_bytes + rhs_tma_tile_bytes
+      txcount = c(lhs_tma_tile_bytes + rhs_tma_tile_bytes)
       common_copy_args = dict(
           swizzle=128, barrier=barrier, arrive=False, uniform=False,
       )
       with single_thread():
-        barrier.arrive_expect_tx(txcount)
+        nvgpu.mbarrier_arrive_expect_tx(barrier_group.value, txcount, slot)
         ctx.async_copy(
             src_ref=a_device,
             dst_ref=memref_slice(compute_smem["lhs"], slot),
@@ -347,7 +349,7 @@ def build_kernel(
       si = arith.remui(ki, c(stages))
 
       with ctx.named_region("TMA wait"):
-        barriers[si].wait()
+        barrier_group[si].wait()
 
       with ctx.named_region("WGMMA"):
         a_slice = memref_slice(compute_smem["lhs"], si)
@@ -437,7 +439,7 @@ def build_kernel(
           jax.ShapeDtypeStruct((n, k) if rhs_transpose else (k, n), rhs_dtype),
       ),
       jax.ShapeDtypeStruct((m, n), jnp.float32),
-      (smem_shape, TMABarrier(stages)),
+      smem_shape,
       profiler_spec,
   )
 
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
index 8b13a00bc..1decd8aca 100644
--- a/jax/experimental/mosaic/gpu/fragmented_array.py
+++ b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -15,7 +15,6 @@
 """Utilities for code generator."""
 
 import dataclasses
-from typing import Callable
 
 import jax
 from jaxlib.mlir import ir
@@ -316,66 +315,30 @@ class FragmentedArray:
   def exp(self, approx: bool = False):
     if not ir.FloatType.isinstance(self.mlir_dtype):
       raise NotImplementedError
-    if approx:
+    def fast_exp(x):
       f32 = ir.F32Type.get()
       if self.mlir_dtype != f32:
         raise NotImplementedError
       log2e = arith.constant(f32, ir.FloatAttr.get(f32, 1.4426950408889634))
-      def fast_exp(x):
-        scaled = arith.mulf(x, log2e)
-        return llvm.inline_asm(f32, [scaled], "ex2.approx.f32 $0, $1;", "=f,f")
-      return self._pointwise(self._lift_fast_unary(fast_exp))
-    return self._pointwise(mlir_math.exp)
-
-  def sin(self, approx: bool = False):
-    if not ir.FloatType.isinstance(self.mlir_dtype):
-      raise NotImplementedError
-    if approx and self.mlir_dtype != ir.F32Type.get():
-      raise NotImplementedError
-    return self._pointwise(
-        self._lift_fast_unary("sin.approx.f32") if approx else mlir_math.sin
-    )
-
-  def cos(self, approx: bool = False):
-    if not ir.FloatType.isinstance(self.mlir_dtype):
-      raise NotImplementedError
-    if approx and self.mlir_dtype != ir.F32Type.get():
-      raise NotImplementedError
-    return self._pointwise(
-        self._lift_fast_unary("cos.approx.f32") if approx else mlir_math.cos
-    )
-
-  def rsqrt(self, approx: bool = False):
-    if not ir.FloatType.isinstance(self.mlir_dtype):
-      raise NotImplementedError
-    if approx and self.mlir_dtype != ir.F32Type.get():
-      raise NotImplementedError
-    return self._pointwise(
-        self._lift_fast_unary("rsqrt.approx.f32") if approx else mlir_math.rsqrt
-    )
-
-  @staticmethod
-  def _lift_fast_unary(
-      instr: str | Callable[[ir.Value], ir.Value],
-  ) -> Callable[[ir.Value], ir.Value]:
-    def fast_instr(x):
-      f32 = ir.F32Type.get()
       if x.type == f32:
-        if isinstance(instr, str):
-          return llvm.inline_asm(f32, [x], instr + " $0, $1;", "=f,f")
-        else:
-          return instr(x)
+        scaled = arith.mulf(x, log2e)
+        return llvm.inline_asm(
+            f32, [scaled], "ex2.approx.f32 $0,$1;", "=f,f", asm_dialect=0
+        )
       elif ir.VectorType.isinstance(x.type):
         index = ir.IndexType.get()
         result = llvm.mlir_undef(x.type)
         for i in range(2):
           v = vector.extractelement(x, position=c(i, index))
-          vr = fast_instr(v)
+          vr = fast_exp(v)
           result = vector.insertelement(vr, result, position=c(i, index))
         return result
       else:
         raise NotImplementedError(x.type)
-    return fast_instr
+    return self._pointwise(fast_exp if approx else mlir_math.exp)
+
+  def rsqrt(self):
+    return self._pointwise(mlir_math.rsqrt)
 
   def __and__(self, other):
     if not ir.IntegerType.isinstance(self.mlir_dtype):
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index ee318df57..32c253255 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -501,91 +501,89 @@ def commit_shared():
   )
 
 
-@dataclasses.dataclass(frozen=True)
-class BarrierRef:
-  base_address: ir.Value
-  offset: ir.Value
-  phases: ir.Value
-  num_barriers: int
+class BarrierArray:
 
-  @staticmethod
-  def initialize(address: ir.Value, num_barriers: int, arrival_count: int = 1) -> "BarrierRef":
+  def __init__(self, num_barriers: int, arrival_count: int = 1):
+    barrier_group_ty = ir.Type.parse(
+        "!nvgpu.mbarrier.group<memorySpace=#gpu.address_space<workgroup>,"
+        f" num_barriers={num_barriers}>"
+    )
+
+    self.num_barriers = num_barriers
+    self.value = nvgpu.mbarrier_create(barrier_group_ty)
+    self.num_barriers = num_barriers
+    index = ir.IndexType.get()
     if num_barriers > 32:
       raise NotImplementedError("Only up to 32 barriers per group supported")
     i32 = ir.IntegerType.get_signless(32)
-    i64 = ir.IntegerType.get_signless(64)
-    ptr = ir.Type.parse("!llvm.ptr<3>")
-    phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
-    memref.store(c(0, i32), phases, [])
+    self.phases = memref.alloca(ir.MemRefType.get((), i32), [], [])
+    memref.store(c(0, i32), self.phases, [])
     with single_thread(per_block=True):
       for i in range(num_barriers):
-        nvvm.mbarrier_init_shared(
-            llvm.getelementptr(ptr, address, [], [i], i64),
-            c(arrival_count, i32),
-        )
-    return BarrierRef(address, c(0, i32), phases, num_barriers)
-
-  def __iter__(self) -> Iterator["BarrierRef"]:
-    if self.num_barriers == 1:
-      raise ValueError("Cannot iterate over a single barrier")
+        nvgpu.mbarrier_init(self.value, c(arrival_count, index), c(i, index))
+    gpu.barrier()
+
+  def __iter__(self) -> Iterator["Barrier"]:
     for offset in range(self.num_barriers):
       yield self[offset]
 
-  def __getitem__(self, offset: ir.Value | int) -> "BarrierRef":
-    if self.num_barriers == 1:
-      raise ValueError("Cannot index a single barrier")
-    i32 = ir.IntegerType.get_signless(32)
+  def __getitem__(self, offset: ir.Value | int):
+    index = ir.IndexType.get()
     if isinstance(offset, int):
-      offset = c(offset, i32)
-    elif ir.IndexType.isinstance(offset.type):
-      offset = arith.index_castui(i32, offset)
-    elif offset.type != i32:
-      raise ValueError(f"Expected a dynamic index or an integer, got {offset}")
-    return BarrierRef(
-        self.base_address,
-        arith.addi(self.offset, offset),
-        self.phases,
-        1,
-    )
+      offset = c(offset, index)
+    if ir.IntegerType.isinstance(offset.type):
+      offset = arith.index_castui(index, offset)
+    return Barrier(self, offset)
+
+
+@dataclasses.dataclass(frozen=True)
+class Barrier:
+  barrier_array: BarrierArray
+  offset: ir.Value
 
   def wait_parity(self, parity, expect_wait=False):
     i1 = ir.IntegerType.get_signless(1)
-    i32 = ir.IntegerType.get_signless(32)
-    ticks = c(10000000, i32)
-    address = self.get_ptr()
-    parity = arith.extui(i32, parity)
+    index = ir.IndexType.get()
     if expect_wait:
-      nvvm.mbarrier_try_wait_parity_shared(address, parity, ticks)
+      nvgpu.mbarrier_try_wait_parity(
+          self.barrier_array.value, parity, c(10000000, index), self.offset,
+      )
+      return
+    barrier_ptr = self.get_ptr()
     barrier_ready = llvm.inline_asm(
         i1,
-        [address, parity],
+        [barrier_ptr, parity],
         "mbarrier.test_wait.parity.shared.b64 $0, [$1], $2;",
         "=b,l,r",
+        asm_dialect=0,
         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)
+      nvgpu.mbarrier_try_wait_parity(
+          self.barrier_array.value, parity, c(10000000, index), self.offset,
+      )
       scf.yield_([])
 
   def wait(self, expect_wait=False):
-    parities = memref.load(self.phases, [])
+    parities = memref.load(self.barrier_array.phases, [])
     parity, new_parities = self.update_parities(parities)
-    memref.store(new_parities, self.phases, [])
+    memref.store(new_parities, self.barrier_array.phases, [])
     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)
-    bitmask = arith.shli(c(1, i32), self.offset)
+    offset_i32 = arith.index_castui(i32, self.offset)
+    bitmask = arith.shli(c(1, i32), offset_i32)
     parity = arith.cmpi(
         arith.CmpIPredicate.ne, arith.andi(parities, bitmask), c(0, i32)
     )
     return parity, arith.xori(parities, bitmask)
 
   def arrive(self):
-    i64 = ir.IntegerType.get_signless(64)
-    nvvm.mbarrier_arrive_shared(i64, self.get_ptr())
+    token_ty = ir.Type.parse("!nvgpu.mbarrier.token")
+    nvgpu.mbarrier_arrive(token_ty, self.barrier_array.value, self.offset)
 
   def arrive_expect_tx(self, bytes: int | ir.Value):
     if isinstance(bytes, int):
@@ -593,53 +591,54 @@ class BarrierRef:
     nvvm.mbarrier_arrive_expect_tx(self.get_ptr(), bytes)
 
   def get_ptr(self):
-    ptr = ir.Type.parse("!llvm.ptr<3>")
+    i32 = ir.IntegerType.get_signless(32)
     i64 = ir.IntegerType.get_signless(64)
-    DYNAMIC32 = -2147483648
+    ptr_ty = ir.Type.parse("!llvm.ptr<3>")
+    smem = ir.IntegerAttr.get(i64, 3)
+    num_barriers = self.barrier_array.num_barriers
+    mbarrier_ref_ty = ir.MemRefType.get((num_barriers,), i64, memory_space=smem)
+    mbarrier_ref = builtin.unrealized_conversion_cast(
+        [mbarrier_ref_ty], [self.barrier_array.value],
+    )
+    mbarrier_ref_ptr = memref.extract_aligned_pointer_as_index(mbarrier_ref)
+    barrier_arr_ptr = llvm.inttoptr(
+        ptr_ty, arith.index_cast(i64, mbarrier_ref_ptr),
+    )
+    offset_i32 = arith.index_cast(i32, self.offset)
     return llvm.getelementptr(
-        ptr, self.base_address, [self.offset], [DYNAMIC32], i64
+        ptr_ty, barrier_arr_ptr, [offset_i32], [-2147483648], i64,
     )
 
 
-@dataclasses.dataclass(frozen=True)
-class CollectiveBarrierRef:
-  barrier: BarrierRef
-  cluster_mask: ir.Value
-
-  @staticmethod
-  def initialize(
-      address: ir.Value,
-      num_barriers: int,
-      dims: Sequence[gpu.Dimension],
-      cluster_shape: tuple[int, int, int],
-  ) -> "CollectiveBarrierRef":
+class CollectiveBarrierArray(BarrierArray):
+  def __init__(self, ctx, dims: Sequence[gpu.Dimension], num_barriers):
     i32 = ir.IntegerType.get_signless(32)
+    self.ctx = ctx
+    self.dims = dims
+    self.cluster_mask = c(0, i32)
     # With the exception of the current device, each pair of slices along
     # collective dims is disjoint. Since the current device is overcounted,
     # we must decrease the arrival count a little.
-    arrival_count = sum(cluster_shape[d] for d in dims) - len(dims) + 1
-    cluster_mask = c(0, i32)
+    arrival_count = sum(ctx.cluster_size[d] for d in dims) - len(dims) + 1
+    super().__init__(num_barriers, arrival_count=arrival_count)
     for d in dims:
-      cluster_mask = arith.ori(
-          cluster_mask, cluster_collective_mask(cluster_shape, d)
+      self.cluster_mask = arith.ori(
+          self.cluster_mask, ctx.cluster_collective_mask(d)
       )
-    barrier = BarrierRef.initialize(address, num_barriers, arrival_count=arrival_count)
-    return CollectiveBarrierRef(barrier, cluster_mask)
-
-  def __iter__(self, offset):
-    for b in self.barrier:
-      yield CollectiveBarrierRef(b, self.cluster_mask)
 
   def __getitem__(self, offset):
-    return CollectiveBarrierRef(self.barrier[offset], self.cluster_mask)
+    return CollectiveBarrier(super().__getitem__(offset))
+
+
+@dataclasses.dataclass(frozen=True)
+class CollectiveBarrier:
+  barrier: Barrier
 
   def arrive(self):
     """Arrives on a barrier in all blocks that share at least one of the coordinates along the collective dimensions.
 
     Note that unlike in arrive, each warpgroup arrives once.
     """
-    if self.barrier.num_barriers != 1:
-      raise ValueError("Can only arrive on a single barrier")
     i32 = ir.IntegerType.get_signless(32)
     thread_in_warpgroup = arith.remui(thread_idx(), c(WARPGROUP_SIZE, i32))
     signaled_block = arith.divui(
@@ -647,7 +646,10 @@ class CollectiveBarrierRef:
     )
     is_collective_block = arith.cmpi(
         arith.CmpIPredicate.ne,
-        arith.andi(self.cluster_mask, arith.shli(c(1, i32), signaled_block)),
+        arith.andi(
+            self.barrier.barrier_array.cluster_mask,
+            arith.shli(c(1, i32), signaled_block),
+        ),
         c(0, i32),
     )
     is_signaling_thread = arith.cmpi(
@@ -860,33 +862,3 @@ def memref_ptr(memref_arg, memory_space=None):
           overflow_flags=llvm.IntegerOverflowFlags.none,
       ),
   )
-
-
-def cluster_collective_mask(
-    cluster_shape: tuple[int, int, int], collective: gpu.Dimension
-):
-  # We first compute the linearized index of the slice along the collective
-  # dim that contains the current block. Then, the mask is a sequence of 1s
-  # strided by the position of the collective dim, shifted left by the linear
-  # slice index.
-  # TODO(apaszke): Make sure this gets hoisted outside of any loops.
-  # If not, we might need to do it manually.
-  i32 = ir.IntegerType.get_signless(32)
-  stride = 1
-  mask_shift = c(0, i32)
-  collective_stride = None
-  for cluster_dim in gpu.Dimension:
-    if cluster_shape[cluster_dim] == 1:
-      continue
-    if cluster_dim != collective:
-      dim_idx = arith.index_castui(i32, gpu.cluster_block_id(cluster_dim))
-      mask_shift = arith.addi(
-          mask_shift, arith.muli(dim_idx, c(stride, i32)),
-      )
-    else:
-      collective_stride = stride
-    stride *= cluster_shape[cluster_dim]
-  mask_unshifted = 0
-  for i in range(cluster_shape[collective]):
-    mask_unshifted |= 1 << (i * collective_stride)
-  return arith.shli(c(mask_unshifted, i32), mask_shift)
diff --git a/jaxlib/cpu/cpu_kernels.cc b/jaxlib/cpu/cpu_kernels.cc
index c7bc9dc4b..63bd7ac07 100644
--- a/jaxlib/cpu/cpu_kernels.cc
+++ b/jaxlib/cpu/cpu_kernels.cc
@@ -118,10 +118,6 @@ XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(
 
 // FFI Kernels
 
-JAX_CPU_REGISTER_HANDLER(blas_strsm_ffi);
-JAX_CPU_REGISTER_HANDLER(blas_dtrsm_ffi);
-JAX_CPU_REGISTER_HANDLER(blas_ctrsm_ffi);
-JAX_CPU_REGISTER_HANDLER(blas_ztrsm_ffi);
 JAX_CPU_REGISTER_HANDLER(lapack_sgetrf_ffi);
 JAX_CPU_REGISTER_HANDLER(lapack_dgetrf_ffi);
 JAX_CPU_REGISTER_HANDLER(lapack_cgetrf_ffi);
diff --git a/jaxlib/cpu/lapack.cc b/jaxlib/cpu/lapack.cc
index 0aa924486..602e4dd4e 100644
--- a/jaxlib/cpu/lapack.cc
+++ b/jaxlib/cpu/lapack.cc
@@ -69,10 +69,6 @@ void GetLapackKernelsFromScipy() {
   AssignKernelFn<Trsm<double>>(blas_ptr("dtrsm"));
   AssignKernelFn<Trsm<std::complex<float>>>(blas_ptr("ctrsm"));
   AssignKernelFn<Trsm<std::complex<double>>>(blas_ptr("ztrsm"));
-  AssignKernelFn<TriMatrixEquationSolver<DataType::F32>>(blas_ptr("strsm"));
-  AssignKernelFn<TriMatrixEquationSolver<DataType::F64>>(blas_ptr("dtrsm"));
-  AssignKernelFn<TriMatrixEquationSolver<DataType::C64>>(blas_ptr("ctrsm"));
-  AssignKernelFn<TriMatrixEquationSolver<DataType::C128>>(blas_ptr("ztrsm"));
 
   nb::module_ cython_lapack =
       nb::module_::import_("scipy.linalg.cython_lapack");
@@ -223,10 +219,6 @@ nb::dict Registrations() {
   dict["lapack_zhetrd"] =
       EncapsulateFunction(Sytrd<std::complex<double>>::Kernel);
 
-  dict["blas_strsm_ffi"] = EncapsulateFunction(blas_strsm_ffi);
-  dict["blas_dtrsm_ffi"] = EncapsulateFunction(blas_dtrsm_ffi);
-  dict["blas_ctrsm_ffi"] = EncapsulateFunction(blas_ctrsm_ffi);
-  dict["blas_ztrsm_ffi"] = EncapsulateFunction(blas_ztrsm_ffi);
   dict["lapack_sgetrf_ffi"] = EncapsulateFunction(lapack_sgetrf_ffi);
   dict["lapack_dgetrf_ffi"] = EncapsulateFunction(lapack_dgetrf_ffi);
   dict["lapack_cgetrf_ffi"] = EncapsulateFunction(lapack_cgetrf_ffi);
diff --git a/jaxlib/cpu/lapack.h b/jaxlib/cpu/lapack.h
index 81a5acce5..828a147d7 100644
--- a/jaxlib/cpu/lapack.h
+++ b/jaxlib/cpu/lapack.h
@@ -23,18 +23,6 @@ namespace jax {
 
 // FFI Definition Macros (by DataType)
 
-#define JAX_CPU_DEFINE_TRSM(name, data_type)                                  \
-  XLA_FFI_DEFINE_HANDLER(name, TriMatrixEquationSolver<data_type>::Kernel,    \
-                         ::xla::ffi::Ffi::Bind()                              \
-                             .Arg<::xla::ffi::Buffer<data_type>>(/*x*/)       \
-                             .Arg<::xla::ffi::Buffer<data_type>>(/*y*/)       \
-                             .Arg<::xla::ffi::BufferR0<data_type>>(/*alpha*/) \
-                             .Ret<::xla::ffi::Buffer<data_type>>(/*y_out*/)   \
-                             .Attr<MatrixParams::Side>("side")                \
-                             .Attr<MatrixParams::UpLo>("uplo")                \
-                             .Attr<MatrixParams::Transpose>("trans_x")        \
-                             .Attr<MatrixParams::Diag>("diag"))
-
 #define JAX_CPU_DEFINE_GETRF(name, data_type)                \
   XLA_FFI_DEFINE_HANDLER(                                    \
       name, LuDecomposition<data_type>::Kernel,              \
@@ -104,11 +92,6 @@ namespace jax {
 
 // FFI Handlers
 
-JAX_CPU_DEFINE_TRSM(blas_strsm_ffi, ::xla::ffi::DataType::F32);
-JAX_CPU_DEFINE_TRSM(blas_dtrsm_ffi, ::xla::ffi::DataType::F64);
-JAX_CPU_DEFINE_TRSM(blas_ctrsm_ffi, ::xla::ffi::DataType::C64);
-JAX_CPU_DEFINE_TRSM(blas_ztrsm_ffi, ::xla::ffi::DataType::C128);
-
 JAX_CPU_DEFINE_GETRF(lapack_sgetrf_ffi, ::xla::ffi::DataType::F32);
 JAX_CPU_DEFINE_GETRF(lapack_dgetrf_ffi, ::xla::ffi::DataType::F64);
 JAX_CPU_DEFINE_GETRF(lapack_cgetrf_ffi, ::xla::ffi::DataType::C64);
@@ -134,7 +117,6 @@ JAX_CPU_DEFINE_GESDD(lapack_dgesdd_ffi, ::xla::ffi::DataType::F64);
 JAX_CPU_DEFINE_GESDD_COMPLEX(lapack_cgesdd_ffi, ::xla::ffi::DataType::C64);
 JAX_CPU_DEFINE_GESDD_COMPLEX(lapack_zgesdd_ffi, ::xla::ffi::DataType::C128);
 
-#undef JAX_CPU_DEFINE_TRSM
 #undef JAX_CPU_DEFINE_GETRF
 #undef JAX_CPU_DEFINE_GEQRF
 #undef JAX_CPU_DEFINE_ORGQR
diff --git a/jaxlib/cpu/lapack_kernels.cc b/jaxlib/cpu/lapack_kernels.cc
index 85fa02893..08a0537eb 100644
--- a/jaxlib/cpu/lapack_kernels.cc
+++ b/jaxlib/cpu/lapack_kernels.cc
@@ -160,48 +160,6 @@ template struct Trsm<double>;
 template struct Trsm<std::complex<float>>;
 template struct Trsm<std::complex<double>>;
 
-// FFI Kernel
-
-template <ffi::DataType dtype>
-ffi::Error TriMatrixEquationSolver<dtype>::Kernel(
-    ffi::Buffer<dtype> x, ffi::Buffer<dtype> y, ffi::BufferR0<dtype> alpha,
-    ffi::ResultBuffer<dtype> y_out, MatrixParams::Side side,
-    MatrixParams::UpLo uplo, MatrixParams::Transpose trans_x,
-    MatrixParams::Diag diag) {
-  CopyIfDiffBuffer(y, y_out);
-
-  auto [batch_count, y_rows, y_cols] = SplitBatch2D(y.dimensions());
-  auto* y_out_data = y_out->typed_data();
-  lapack_int x_leading_dim_v =
-      side == MatrixParams::Side::kLeft ? y_rows : y_cols;
-  lapack_int y_leading_dim_v = y_rows;
-
-  auto side_v = static_cast<char>(side);
-  auto uplo_v = static_cast<char>(uplo);
-  auto trans_x_v = static_cast<char>(trans_x);
-  auto diag_v = static_cast<char>(diag);
-  FFI_ASSIGN_OR_RETURN(auto y_rows_v, MaybeCastNoOverflow<lapack_int>(y_rows));
-  FFI_ASSIGN_OR_RETURN(auto y_cols_v, MaybeCastNoOverflow<lapack_int>(y_cols));
-
-  auto* x_data = x.typed_data();
-  const int64_t y_out_step{y_rows * y_cols};
-  const int64_t x_step{x_leading_dim_v * x_leading_dim_v};
-  for (int64_t i = 0; i < batch_count; ++i) {
-    fn(&side_v, &uplo_v, &trans_x_v, &diag_v, &y_rows_v, &y_cols_v,
-       alpha.typed_data(), x_data, &x_leading_dim_v, y_out_data,
-       &y_leading_dim_v);
-
-    y_out_data += y_out_step;
-    x_data += x_step;
-  }
-  return ffi::Error::Success();
-}
-
-template struct TriMatrixEquationSolver<ffi::DataType::F32>;
-template struct TriMatrixEquationSolver<ffi::DataType::F64>;
-template struct TriMatrixEquationSolver<ffi::DataType::C64>;
-template struct TriMatrixEquationSolver<ffi::DataType::C128>;
-
 //== LU Decomposition ==//
 
 // lapack getrf
diff --git a/jaxlib/cpu/lapack_kernels.h b/jaxlib/cpu/lapack_kernels.h
index fd9d8c975..ea67bd32c 100644
--- a/jaxlib/cpu/lapack_kernels.h
+++ b/jaxlib/cpu/lapack_kernels.h
@@ -109,24 +109,6 @@ struct Trsm {
   static void Kernel(void* out, void** data, XlaCustomCallStatus*);
 };
 
-// FFI Kernel
-
-template <::xla::ffi::DataType dtype>
-struct TriMatrixEquationSolver {
-  using ValueType = ::xla::ffi::NativeType<dtype>;
-  using FnType = void(char* side, char* uplo, char* transa, char* diag,
-                      lapack_int* m, lapack_int* n, ValueType* alpha,
-                      ValueType* a, lapack_int* lda, ValueType* b,
-                      lapack_int* ldb);
-
-  inline static FnType* fn = nullptr;
-  static ::xla::ffi::Error Kernel(
-      ::xla::ffi::Buffer<dtype> x, ::xla::ffi::Buffer<dtype> y,
-      ::xla::ffi::BufferR0<dtype> alpha, ::xla::ffi::ResultBuffer<dtype> y_out,
-      MatrixParams::Side side, MatrixParams::UpLo uplo,
-      MatrixParams::Transpose trans_x, MatrixParams::Diag diag);
-};
-
 //== LU Decomposition ==//
 
 // lapack getrf
diff --git a/jaxlib/cpu/lapack_kernels_using_lapack.cc b/jaxlib/cpu/lapack_kernels_using_lapack.cc
index d17925e13..9c8a7b61e 100644
--- a/jaxlib/cpu/lapack_kernels_using_lapack.cc
+++ b/jaxlib/cpu/lapack_kernels_using_lapack.cc
@@ -26,10 +26,10 @@ namespace ffi = xla::ffi;
 
 extern "C" {
 
-jax::TriMatrixEquationSolver<ffi::DataType::F32>::FnType strsm_;
-jax::TriMatrixEquationSolver<ffi::DataType::F64>::FnType dtrsm_;
-jax::TriMatrixEquationSolver<ffi::DataType::C64>::FnType ctrsm_;
-jax::TriMatrixEquationSolver<ffi::DataType::C128>::FnType ztrsm_;
+jax::Trsm<float>::FnType strsm_;
+jax::Trsm<double>::FnType dtrsm_;
+jax::Trsm<std::complex<float>>::FnType ctrsm_;
+jax::Trsm<std::complex<double>>::FnType ztrsm_;
 
 jax::LuDecomposition<ffi::DataType::F32>::FnType sgetrf_;
 jax::LuDecomposition<ffi::DataType::F64>::FnType dgetrf_;
@@ -87,22 +87,6 @@ namespace jax {
 
 #define JAX_KERNEL_FNTYPE_MISMATCH_MSG "FFI Kernel FnType mismatch"
 
-static_assert(
-    std::is_same_v<jax::TriMatrixEquationSolver<ffi::DataType::F32>::FnType,
-                   jax::Trsm<float>::FnType>,
-    JAX_KERNEL_FNTYPE_MISMATCH_MSG);
-static_assert(
-    std::is_same_v<jax::TriMatrixEquationSolver<ffi::DataType::F64>::FnType,
-                   jax::Trsm<double>::FnType>,
-    JAX_KERNEL_FNTYPE_MISMATCH_MSG);
-static_assert(
-    std::is_same_v<jax::TriMatrixEquationSolver<ffi::DataType::C64>::FnType,
-                   jax::Trsm<std::complex<float>>::FnType>,
-    JAX_KERNEL_FNTYPE_MISMATCH_MSG);
-static_assert(
-    std::is_same_v<jax::TriMatrixEquationSolver<ffi::DataType::C128>::FnType,
-                   jax::Trsm<std::complex<double>>::FnType>,
-    JAX_KERNEL_FNTYPE_MISMATCH_MSG);
 static_assert(std::is_same_v<jax::LuDecomposition<ffi::DataType::F32>::FnType,
                              jax::Getrf<float>::FnType>,
               JAX_KERNEL_FNTYPE_MISMATCH_MSG);
@@ -234,11 +218,6 @@ static auto init = []() -> int {
 
   // FFI Kernels
 
-  AssignKernelFn<TriMatrixEquationSolver<ffi::DataType::F32>>(strsm_);
-  AssignKernelFn<TriMatrixEquationSolver<ffi::DataType::F64>>(dtrsm_);
-  AssignKernelFn<TriMatrixEquationSolver<ffi::DataType::C64>>(ctrsm_);
-  AssignKernelFn<TriMatrixEquationSolver<ffi::DataType::C128>>(ztrsm_);
-
   AssignKernelFn<LuDecomposition<ffi::DataType::F32>>(sgetrf_);
   AssignKernelFn<LuDecomposition<ffi::DataType::F64>>(dgetrf_);
   AssignKernelFn<LuDecomposition<ffi::DataType::C64>>(cgetrf_);


Test patch:
diff --git a/tests/export_back_compat_test.py b/tests/export_back_compat_test.py
index a8b856001..d510aab9d 100644
--- a/tests/export_back_compat_test.py
+++ b/tests/export_back_compat_test.py
@@ -479,6 +479,7 @@ class CompatTest(bctu.CompatTestBase):
                                   np.asarray(out), atol=1e-4, rtol=1e-4))
 
   @jtu.parameterized_filterable(
+    one_containing="f32",
     kwargs=[
       dict(testcase_name=f"_dtype={dtype_name}", dtype_name=dtype_name)
       for dtype_name in ("f32", "f64", "c64", "c128")])
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index 14d4f4277..45c948abd 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -542,7 +542,7 @@ class WGMMATest(TestCase):
     rhs_order = col_major if rhs_transpose else row_major
 
     def kernel(ctx, lhs, rhs, out, scratch):
-      lhs_smem, rhs_smem, barriers = scratch
+      lhs_smem, rhs_smem = scratch
       if tma_inputs:
         lhs_transform = (mosaic_gpu.TileTransform((64, nk_tile)),)
         if lhs_transpose:
@@ -551,6 +551,7 @@ class WGMMATest(TestCase):
         rhs_transform = (mosaic_gpu.TileTransform((nk_tile, nk_tile)),)
         if rhs_transpose:
           rhs_transform += (mosaic_gpu.TransposeTransform((1, 0, 2, 3)),)
+        barriers = BarrierArray(2)
         ctx.async_copy(
             src_ref=lhs,
             dst_ref=lhs_smem,
@@ -615,7 +616,6 @@ class WGMMATest(TestCase):
         jax.ShapeDtypeStruct(
             (k // nk_tile, n // nk_tile, nk_tile, nk_tile), in_jax_dtype
         ),
-        mgpu.TMABarrier(2),
     ]
     z = mosaic_gpu.as_gpu_kernel(
         kernel, (1, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape
@@ -690,9 +690,10 @@ class BarrierTest(TestCase):
 
   def test_wg_communication(self):
     i32 = ir.IntegerType.get_signless(32)
-    def kernel(ctx, dst, scratch):
-      tmp, barriers = scratch
+    def kernel(ctx, dst, tmp):
       del ctx  # Unused.
+      barriers = BarrierArray(3, arrival_count=128)
+      gpu.barrier()  # Make sure the barriers are initialized.
       wg_idx = arith.divui(mgpu.warp_idx(), c(4, i32))
       is_first_wg = arith.cmpi(arith.CmpIPredicate.eq, wg_idx, c(0, i32))
       is_second_wg = arith.cmpi(arith.CmpIPredicate.eq, wg_idx, c(1, i32))
@@ -724,10 +725,7 @@ class BarrierTest(TestCase):
         (2 * 128, 1, 1),
         (),
         out_shape,
-        (
-            jax.ShapeDtypeStruct((128,), jnp.int32),
-            mgpu.Barrier(arrival_count=128, num_barriers=3),
-        ),
+        jax.ShapeDtypeStruct((128,), jnp.int32),
     )()
     np.testing.assert_array_equal(y, np.full_like(y, 3, dtype=np.int32))
 
@@ -751,7 +749,11 @@ class BarrierTest(TestCase):
     cluster = [1, 1, 1]
     for d in itertools.chain(collective_dims, noncollective_dims):
       cluster[d] = 2
-    def kernel(ctx, dst, collective_barrier):
+    def kernel(ctx, dst, _):
+      collective_barrier = CollectiveBarrierArray(ctx, collective_dims, 1)[0]
+      nvvm.fence_mbarrier_init()
+      nvvm.cluster_arrive_relaxed()
+      nvvm.cluster_wait()
       collective_barrier.arrive()
       collective_barrier.wait()
       tid = thread_idx()
@@ -762,9 +764,8 @@ class BarrierTest(TestCase):
         stride = arith.muli(stride, gpu.grid_dim(d))
       memref.store(arith.index_cast(i32, linear_idx), dst, [linear_idx])
     out_shape = jax.ShapeDtypeStruct((math.prod(cluster) * 128,), jnp.int32)
-    scratch = mgpu.ClusterBarrier(collective_dims)
     y = mosaic_gpu.as_gpu_kernel(
-        kernel, cluster, (128, 1, 1), (), out_shape, scratch, cluster=cluster,
+        kernel, cluster, (128, 1, 1), (), out_shape, (), cluster=cluster,
     )()
     np.testing.assert_array_equal(
         y, np.arange(math.prod(cluster) * 128, dtype=np.int32)
@@ -782,14 +783,13 @@ class TMATest(TestCase):
     minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
     shape = (*shape[:-1], minor_size)
     i1 = ir.IntegerType.get_signless(1)
-    def kernel(ctx, src, dst, smem):
-      tmp, barrier = smem
+    def kernel(ctx, src, dst, tmp):
+      barrier = BarrierArray(1)[0]
       ctx.async_copy(src_ref=src, dst_ref=tmp, swizzle=swizzle, barrier=barrier)
       barrier.wait_parity(c(0, i1))
       copy(tmp, dst, swizzle=swizzle)
     x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
-    smem = (x, mgpu.TMABarrier())
-    y = mosaic_gpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
+    y = mosaic_gpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)
     np.testing.assert_array_equal(y, x)
 
   @parameterized.named_parameters(
@@ -818,8 +818,7 @@ class TMATest(TestCase):
     # Note that this kernel does not use the non-collective dimensions in any
     # interesting way and so they don't really have to be part of the cluster.
     # We use them to test that the multicast mask is generated correctly.
-    def kernel(ctx, src, dst, scratch):
-      tmp, barrier = scratch
+    def kernel(ctx, src, dst, tmp):
       stride = 1
       noncollective_idx = c(0, index)
       for d in noncollective_dims:
@@ -828,6 +827,10 @@ class TMATest(TestCase):
             arith.muli(gpu.cluster_block_id(d), c(stride, index))
         )
         stride *= 2
+      barrier = BarrierArray(1)[0]
+      nvvm.fence_mbarrier_init()
+      nvvm.cluster_arrive_relaxed()
+      nvvm.cluster_wait()
       ctx.async_copy(
           src_ref=src,
           dst_ref=tmp,
@@ -846,7 +849,7 @@ class TMATest(TestCase):
           swizzle=swizzle,
       )
     x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
-    smem_shape = (jax.ShapeDtypeStruct(shape[1:], dtype), mgpu.TMABarrier())
+    smem_shape = jax.ShapeDtypeStruct(shape[1:], dtype)
     y = mosaic_gpu.as_gpu_kernel(
         kernel, cluster, (128, 1, 1), x, x, smem_shape, cluster=cluster
     )(x)
@@ -864,8 +867,8 @@ class TMATest(TestCase):
     index = ir.IndexType.get()
     tiling = (32, (swizzle or 128) // jnp.dtype(dtype).itemsize)
     tiled_shape = tile_shape(shape, tiling)[:len(shape)]
-    def kernel(ctx, src, dst, scratch):
-      tmp, barrier = scratch
+    def kernel(ctx, src, dst, tmp):
+      barrier = BarrierArray(1)[0]
       ctx.async_copy(
           src_ref=src,
           dst_ref=tmp,
@@ -882,10 +885,7 @@ class TMATest(TestCase):
         )
         copy(memref_slice(tmp, idxs), memref_slice(dst, s), swizzle=swizzle)
     x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
-    smem = (
-        jax.ShapeDtypeStruct(tile_shape(shape, tiling), dtype),
-        mgpu.TMABarrier(),
-    )
+    smem = jax.ShapeDtypeStruct(tile_shape(shape, tiling), dtype)
     f = mosaic_gpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)
     y = f(x)
     np.testing.assert_array_equal(y, x)
@@ -899,8 +899,8 @@ class TMATest(TestCase):
     # swizzle 32 and 64.
     minor_size = 64 if swizzle is None else swizzle // jnp.dtype(dtype).itemsize
     shape = (4, 5, minor_size)
-    def kernel(ctx, src, dst, smem):
-      tmp, barrier = smem
+    def kernel(ctx, src, dst, tmp):
+      barrier = BarrierArray(1)[0]
       for i in range(4):
         ctx.async_copy(
             src_ref=src,
@@ -912,15 +912,14 @@ class TMATest(TestCase):
         barrier.wait()
       copy(tmp, dst, swizzle=swizzle)
     x = np.arange(np.prod(shape), dtype=dtype).reshape(shape)
-    smem = (x, mgpu.TMABarrier())
-    y = mosaic_gpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, smem)(x)
+    y = mosaic_gpu.as_gpu_kernel(kernel, (1, 1, 1), (128, 1, 1), x, x, x)(x)
     np.testing.assert_array_equal(y, x)
 
   def test_parity_tracking(self):
     shape = (16, 64)
     index = ir.IndexType.get()
-    def kernel(ctx, src, dst, smem):
-      tmp, barrier = smem
+    def kernel(ctx, src, dst, tmp):
+      barrier = BarrierArray(1)[0]
       for i in range(shape[0]):
         s = ds(c(i, index), 1)
         ctx.async_copy(
@@ -930,7 +929,7 @@ class TMATest(TestCase):
         copy(tmp, memref_slice(dst, s))
     x = np.arange(np.prod(shape), dtype=jnp.float16).reshape(shape)
     y = mosaic_gpu.as_gpu_kernel(
-        kernel, (1, 1, 1), (128, 1, 1), x, x, (x[0:1], mgpu.TMABarrier())
+        kernel, (1, 1, 1), (128, 1, 1), x, x, x[0:1]
     )(x)
     np.testing.assert_array_equal(y, x)
 
@@ -988,21 +987,12 @@ class FragmentedArrayTest(TestCase):
         np.testing.assert_array_equal(result, np_op(ref_x, ref_rhs))
 
   @parameterized.product(
-      ops=(
-          (lambda x: mgpu.FragmentedArray.exp(x), np.exp, False),
-          (lambda x: mgpu.FragmentedArray.exp(x, approx=True), np.exp, True),
-          (lambda x: mgpu.FragmentedArray.sin(x), np.sin, False),
-          (lambda x: mgpu.FragmentedArray.sin(x, approx=True), np.sin, True),
-          (lambda x: mgpu.FragmentedArray.cos(x), np.cos, False),
-          (lambda x: mgpu.FragmentedArray.cos(x, approx=True), np.cos, True),
-          (lambda x: mgpu.FragmentedArray.rsqrt(x), jax.lax.rsqrt, False),
-          (lambda x: mgpu.FragmentedArray.rsqrt(x, approx=True), jax.lax.rsqrt, True),
-      ),
+      ops=((lambda x: mgpu.FragmentedArray.exp(x), np.exp),),
       m=(64, 128),
       n=(8, 16, 32, 64, 80, 128, 256),
   )
   def test_unary(self, ops, m=64, n=32):
-    op, np_op, is_approx = ops
+    op, np_op = ops
     def kernel(ctx, dst, _):
       f32 = ir.F32Type.get()
       iota = iota_tensor(m=m, n=n, mlir_dtype=f32)
@@ -1012,9 +1002,7 @@ class FragmentedArrayTest(TestCase):
         kernel, (1, 1, 1), (128, 1, 1), (), out_shape, ()
     )()
     x = np.arange(m * n, dtype=jnp.float32).reshape(m, n)
-    atol = 5e-3 if is_approx else 2e-7
-    rtol = 4e-6 if is_approx else 2e-7
-    np.testing.assert_allclose(result, np_op(x), atol=atol, rtol=rtol)
+    np.testing.assert_allclose(result, np_op(x), atol=2e-7, rtol=2e-7)
 
   @parameterized.product(
       op=(arith.addf, arith.maximumf),
diff --git a/tests/pallas/gpu_ops_test.py b/tests/pallas/gpu_ops_test.py
index e7b7a4daa..66233585d 100644
--- a/tests/pallas/gpu_ops_test.py
+++ b/tests/pallas/gpu_ops_test.py
@@ -26,7 +26,7 @@ from jax import random
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.lax.control_flow.for_loop import for_loop
-from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
 from jax.experimental import pallas as pl
 from jax.experimental.pallas.ops.gpu import attention
 from jax.experimental.pallas.ops.gpu import layer_norm
@@ -129,7 +129,7 @@ class PallasBaseTest(jtu.JaxTestCase):
       self.skipTest("Only works on non-Windows platforms")
 
     super().setUp()
-    _trace_kernel_to_jaxpr.cache_clear()
+    _trace_to_jaxpr.cache_clear()
 
   def pallas_call(self, *args, **kwargs):
     return pl.pallas_call(*args, **kwargs, interpret=self.INTERPRET)
diff --git a/tests/pallas/pallas_shape_poly_test.py b/tests/pallas/pallas_shape_poly_test.py
index e9384afec..a845d5b96 100644
--- a/tests/pallas/pallas_shape_poly_test.py
+++ b/tests/pallas/pallas_shape_poly_test.py
@@ -27,7 +27,7 @@ from absl.testing import parameterized
 import jax
 from jax._src import config
 from jax._src import test_util as jtu
-from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
 import jax.numpy as jnp
 from jax.experimental import pallas as pl
 from jax import export
@@ -93,7 +93,7 @@ class ShapePolyTest(jtu.JaxTestCase,
     if sys.platform == "win32":
       self.skipTest("Only works on non-Windows platforms")
     super().setUp()
-    _trace_kernel_to_jaxpr.cache_clear()
+    _trace_to_jaxpr.cache_clear()
 
   def test_copy(self):
     # The blocks are static, but the input and the grid are of polymorphic
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
index 73b4d3eda..f0c3472a7 100644
--- a/tests/pallas/pallas_test.py
+++ b/tests/pallas/pallas_test.py
@@ -32,7 +32,7 @@ from jax._src import config
 from jax._src import dtypes
 from jax._src import test_util as jtu
 from jax._src.lax.control_flow.for_loop import for_loop
-from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
 from jax.experimental import pallas as pl
 import jax.numpy as jnp
 import numpy as np
@@ -145,13 +145,14 @@ class PallasBaseTest(jtu.JaxTestCase):
       self.skipTest("Only works on non-Windows platforms")
 
     super().setUp()
-    _trace_kernel_to_jaxpr.cache_clear()
+    _trace_to_jaxpr.cache_clear()
 
   def pallas_call(self, *args, **kwargs):
     return pl.pallas_call(*args, **kwargs, interpret=self.INTERPRET)
 
 
 class PallasCallTest(PallasBaseTest):
+
   def test_add_one(self):
     if jtu.test_device_matches(["tpu"]) and not self.INTERPRET:
       self.skipTest("On TPU the test works only in interpret mode")
@@ -654,6 +655,7 @@ class PallasCallInterpreterTest(PallasCallTest):
 
 
 class ApiErrorTest(PallasBaseTest):
+
   def test_pallas_kernel_args_mismatch(self):
     a = np.arange(256, dtype=np.int32)
     f = self.pallas_call(lambda x_ref: None,  # Missing o_ref
@@ -717,7 +719,7 @@ class ApiErrorTest(PallasBaseTest):
                          in_specs=[pl.BlockSpec((4,), lambda: (0, 0))])
     with self.assertRaisesRegex(
         ValueError,
-        "Index map for inputs\\[0\\] must return 1 values to match .*Currently returning 2 values."):
+        "Index map for input\\[0\\] must return 1 values to match .*Currently returning 2 values."):
       f(a)
 
   def test_pallas_call_index_map_captures_consts(self):
@@ -728,7 +730,7 @@ class ApiErrorTest(PallasBaseTest):
                          in_specs=[pl.BlockSpec((4,), lambda: index_map_result)])
     with self.assertRaisesRegex(
         NotImplementedError,
-        "Index map for inputs\\[0\\] captures constants"):
+        "Index map for input\\[0\\] captures constants"):
       f(a)
 
   def test_pallas_call_out_specs_mismatch_shape(self):
@@ -750,7 +752,7 @@ class ApiErrorTest(PallasBaseTest):
                          in_specs=[pl.BlockSpec((1, 1), lambda: (0, 0))])
     with self.assertRaisesRegex(
         ValueError,
-        "Block shape for inputs\\[0\\] .* must have the same number of dimensions as the "
+        "Block shape for input\\[0\\] .* must have the same number of dimensions as the "
         "array shape"):
 
       f(a)
@@ -760,7 +762,7 @@ class ApiErrorTest(PallasBaseTest):
                          out_specs=[pl.BlockSpec((1, 1), lambda: 0)])
     with self.assertRaisesRegex(
         ValueError,
-        "Block shape for outputs\\[0\\] .* must have the same number of dimensions as the "
+        "Block shape for output\\[0\\] .* must have the same number of dimensions as the "
         "array shape"):
       f(a)
 
@@ -1891,6 +1893,7 @@ class PallasCheckifyInterpreterTest(PallasBaseTest):
 
 
 class PallasCallNamedGridTest(PallasBaseTest):
+
   def test_named_grid(self):
 
     def kernel(x_ref, y_ref):
diff --git a/tests/pallas/pallas_vmap_test.py b/tests/pallas/pallas_vmap_test.py
index 8006a7889..cfeff5cf7 100644
--- a/tests/pallas/pallas_vmap_test.py
+++ b/tests/pallas/pallas_vmap_test.py
@@ -24,7 +24,7 @@ from jax import random
 from jax._src.lib import xla_extension
 from jax._src import config
 from jax._src import test_util as jtu
-from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
 from jax.experimental import pallas as pl
 import jax.numpy as jnp
 import numpy as np
@@ -52,7 +52,7 @@ class PallasBaseTest(jtu.JaxTestCase):
       self.skipTest("Only works on non-Windows platforms")
 
     super().setUp()
-    _trace_kernel_to_jaxpr.cache_clear()
+    _trace_to_jaxpr.cache_clear()
 
   def pallas_call(self, *args, **kwargs):
     return pl.pallas_call(*args, **kwargs, interpret=self.INTERPRET)
diff --git a/tests/pallas/tpu_pallas_test.py b/tests/pallas/tpu_pallas_test.py
index 0f0682214..d4586bacb 100644
--- a/tests/pallas/tpu_pallas_test.py
+++ b/tests/pallas/tpu_pallas_test.py
@@ -29,7 +29,7 @@ from jax._src import state
 from jax._src import test_util as jtu
 from jax._src.interpreters import partial_eval as pe
 from jax._src.lib import xla_extension
-from jax._src.pallas.pallas_call import _trace_kernel_to_jaxpr
+from jax._src.pallas.pallas_call import _trace_to_jaxpr
 from jax.experimental import mesh_utils
 from jax.experimental import mosaic
 from jax.experimental import pallas as pl
@@ -65,7 +65,7 @@ class PallasBaseTest(jtu.JaxTestCase):
     if not jtu.test_device_matches(['tpu']) and not self.INTERPRET:
       self.skipTest('Test requires TPUs, or interpret mode')
     super().setUp()
-    _trace_kernel_to_jaxpr.cache_clear()
+    _trace_to_jaxpr.cache_clear()
 
   def pallas_call(self, *args, **kwargs):
     return pl.pallas_call(*args, **kwargs, interpret=self.INTERPRET)
@@ -117,17 +117,12 @@ class PallasCallScalarPrefetchTest(PallasBaseTest):
 
   @jtu.parameterized_filterable(
       kwargs=[
-          dict(scratch=scratch, vmap=vmap, dyn_grid=dyn_grid)
+          dict(scratch=scratch, vmap=vmap)
           for scratch in [True, False]
-          for vmap in [False, True]
-          for dyn_grid in [False, True]
+          for vmap in [True, False]
       ]
   )
-  def test_scalar_prefetch_calling_convention(
-      self, *,
-      scratch: bool, vmap: bool, dyn_grid: bool):
-    # Tests what we process correctly all the various inputs and outputs:
-    # dynamic_grid_dims, index, inputs, outputs, scratch.
+  def test_scalar_prefetch_hoisted_const(self, *, scratch: bool, vmap: bool):
     if jtu.test_device_matches(["cpu"]) and jax.config.x64_enabled:
       self.skipTest("TODO: dslice(start, 1) raises error about slice inputs being int32 and int64")
     # to_store will be hoisted as constants. Choose distinct shapes from in/outs.
@@ -138,39 +133,30 @@ class PallasCallScalarPrefetchTest(PallasBaseTest):
       x_shape = (16, 128)
     x = np.arange(math.prod(x_shape), dtype=np.float32).reshape(x_shape)
 
-    def f(x, grid_size):
+    def f(x):
       s = jnp.array([1, 0], jnp.int32)  # iteration 0 -> 1, iteration 1 -> 0
       @functools.partial(
           self.pallas_call,
           out_shape=jax.ShapeDtypeStruct((64, 128), x.dtype),
           grid_spec=pltpu.PrefetchScalarGridSpec(
-              num_scalar_prefetch=1,  # 1 pytree
-              grid=(grid_size,),
+              num_scalar_prefetch=1,
+              grid=(2,),
               in_specs=[pl.BlockSpec((8, 128),
-                                     lambda i, s_ref: (pl.load(s_ref[0], (i,)), 0))],
+                                     lambda i, s_ref: (pl.load(s_ref, (i,)), 0))],
               out_specs=pl.BlockSpec((32, 128),
-                                     lambda i, s_ref: (pl.load(s_ref[0], i), 0)),
+                                     lambda i, s_ref: (pl.load(s_ref, i), 0)),
               scratch_shapes=([pltpu.SemaphoreType.REGULAR((3,))] if scratch
                               else []),
           ),
       )
-      def kernel(s_refs, src, dst, *scratch_refs):
-        s_ref, s2, s3 = s_refs
-        assert s_ref.shape == (2,)
-        assert s2.shape == (3,)
-        assert s3 is None
+      def kernel(s_ref, src, dst, *scratch_refs):
         store_idx = s_ref[pl.program_id(0)]
         pl.store(dst, (pl.dslice(store_idx, 1), slice(None)), to_store)
-      # Pass a pytree of scalar
-      return kernel((s, np.arange(3, dtype=np.int32), None), x)
+      return kernel(s, x)
 
-    if dyn_grid:
-      f = jax.jit(f)
     if vmap:
-      res = jax.vmap(lambda x: f(x, 2))(x)
-    else:
-      res = f(x, 2)
-
+      f = jax.vmap(f)
+    res = f(x)
     if vmap:
       for i in range(x.shape[0]):
         self.assertAllClose(res[i, 0:1], to_store)
@@ -224,7 +210,7 @@ class PallasCallScalarPrefetchTest(PallasBaseTest):
     x = jnp.ones((16, 128))
     with self.assertRaisesRegex(
         ValueError,
-        r'Index map for inputs\[0\] must return 2 values to match block shape \(8, 128\).'
+        r'Index map for input\[0\] must return 2 values to match block_shape=\(8, 128\).'
         ' Currently returning 1 values.'):
       _ = self.pallas_call(
           body,
@@ -2353,6 +2339,7 @@ class PallasCallTPUBooleanInterpretTest(PallasCallTPUBooleanTest):
 
 
 class PallasCallTPUCheckifyTest(PallasBaseTest):
+
   @parameterized.parameters((2,), (5,), (6,), (7,))
   def test_checkify_with_scalar_prefetch(self, threshold):
     def body(scalar_ref, x_ref, o_ref):
diff --git a/tests/pallas/tpu_splash_attention_kernel_test.py b/tests/pallas/tpu_splash_attention_kernel_test.py
index 3efa259bb..317b6984d 100644
--- a/tests/pallas/tpu_splash_attention_kernel_test.py
+++ b/tests/pallas/tpu_splash_attention_kernel_test.py
@@ -318,6 +318,7 @@ class PallasBaseTest(jtu.JaxTestCase):
 
 
 class SplashAttentionTest(PallasBaseTest):
+
   @parameterized.product(
       is_mqa=(False, True),
       is_segmented=(False, True),
@@ -543,7 +544,7 @@ class SplashAttentionTest(PallasBaseTest):
         data.draw(mha_strategy())
     )
 
-    # Avoid segment ids for rectangular matrices, as it's hard to enforce
+    # Avoid segment ids for rectangular matrices, as its hard to enforce
     # valid masks (non-0 rows).
     hp.assume(q_seq_len == kv_seq_len or not is_segmented)