Repository: google/jax
Pull request #20954 (2024-04-26 17:46:05)

Expose tile function in array_api namespace
Towards https://github.com/google/jax/issues/20200

Exposes `jax.experimental.array_api.tile` which is a thin wrapper around `jax.numpy.tile`


Patch:
diff --git a/docs/jaxpr.rst b/docs/jaxpr.rst
index fae4ee9b5..30075613f 100644
--- a/docs/jaxpr.rst
+++ b/docs/jaxpr.rst
@@ -34,7 +34,7 @@ There are two related representations in the code for jaxprs,
 :py:class:`jax.core.Jaxpr`, and is what you obtain when you use
 :py:func:`jax.make_jaxpr` to inspect jaxprs. It has the following fields:
 
-  * ``jaxpr`` is a :py:class:`jax.core.Jaxpr` representing the actual
+  * ``jaxpr``: is a :py:class:`jax.core.Jaxpr` representing the actual
     computation content of the function (described below).
   * ``consts`` is a list of constants.
 
@@ -42,9 +42,9 @@ The most interesting part of the ClosedJaxpr is the actual execution content,
 represented as a :py:class:`jax.core.Jaxpr` as printed using the following
 grammar::
 
-  Jaxpr ::= { lambda Var* ; Var+. let
-                Eqn*
-              in  [Expr+] }
+   jaxpr ::= { lambda Var* ; Var+.
+               let Eqn*
+               in  [Expr+] }
 
 where:
   * The parameters of the jaxpr are shown as two lists of variables separated by
@@ -62,7 +62,7 @@ where:
 
 Equations are printed as follows::
 
-  Eqn  ::= Var+ = Primitive [ Param* ] Expr+
+  Eqn  ::= let Var+ = Primitive [ Param* ] Expr+
 
 where:
   * ``Var+`` are one or more intermediate variables to be defined as the output
@@ -76,7 +76,7 @@ where:
     square brackets. Each parameter is shown as ``Name = Value``.
 
 
-Most jaxpr primitives are first-order (they take just one or more ``Expr`` as arguments)::
+Most jaxpr primitives are first-order (they take just one or more Expr as arguments)::
 
   Primitive := add | sub | sin | mul | ...
 
diff --git a/jax/_src/pallas/mosaic/__init__.py b/jax/_src/pallas/mosaic/__init__.py
index c8c6c521d..48aea3bd4 100644
--- a/jax/_src/pallas/mosaic/__init__.py
+++ b/jax/_src/pallas/mosaic/__init__.py
@@ -14,26 +14,25 @@
 
 """Module for Mosaic lowering of Pallas call."""
 
-from jax._src.api import named_scope as trace
 from jax._src.pallas.mosaic import core
-from jax._src.pallas.mosaic.core import dma_semaphore
 from jax._src.pallas.mosaic.core import PrefetchScalarGridSpec
-from jax._src.pallas.mosaic.core import semaphore
 from jax._src.pallas.mosaic.core import SemaphoreType
 from jax._src.pallas.mosaic.core import TPUMemorySpace
+from jax._src.pallas.mosaic.core import semaphore
+from jax._src.pallas.mosaic.core import dma_semaphore
 from jax._src.pallas.mosaic.kernel_regeneration_util import encode_kernel_regeneration_metadata
 from jax._src.pallas.mosaic.kernel_regeneration_util import extract_kernel_regeneration_metadata
 from jax._src.pallas.mosaic.lowering import LoweringException
-from jax._src.pallas.mosaic.pipeline import emit_pipeline
 from jax._src.pallas.mosaic.pipeline import emit_pipeline_with_allocations
-from jax._src.pallas.mosaic.pipeline import ManualPrefetchArgs
+from jax._src.pallas.mosaic.pipeline import emit_pipeline
 from jax._src.pallas.mosaic.pipeline import PipelineCallbackArgs
 from jax._src.pallas.mosaic.pipeline import PipelinePrefetchArgs
+from jax._src.pallas.mosaic.pipeline import ManualPrefetchArgs
+from jax._src.pallas.mosaic.primitives import DeviceIdType
 from jax._src.pallas.mosaic.primitives import async_copy
 from jax._src.pallas.mosaic.primitives import async_remote_copy
 from jax._src.pallas.mosaic.primitives import bitcast
 from jax._src.pallas.mosaic.primitives import device_id
-from jax._src.pallas.mosaic.primitives import DeviceIdType
 from jax._src.pallas.mosaic.primitives import get_barrier_semaphore
 from jax._src.pallas.mosaic.primitives import make_async_copy
 from jax._src.pallas.mosaic.primitives import make_async_remote_copy
@@ -43,6 +42,7 @@ from jax._src.pallas.mosaic.primitives import run_scoped
 from jax._src.pallas.mosaic.primitives import semaphore_read
 from jax._src.pallas.mosaic.primitives import semaphore_signal
 from jax._src.pallas.mosaic.primitives import semaphore_wait
+from jax._src.pallas.mosaic.primitives import trace
 
 ANY = TPUMemorySpace.ANY
 CMEM = TPUMemorySpace.CMEM
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index 6e5b2e8ba..04a360f09 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -572,29 +572,6 @@ class LoweringException(Exception):
   pass
 
 
-def _compute_name_stack_updates(
-    old_name_stack: list[str],
-    new_name_stack: list[str]
-) -> tuple[list[str], list[str]]:
-  """Computes the popped/pushed items to the name stack after an update.
-
-  Args:
-    old_name_stack: The name stack prior to the update.
-    new_name_stack: The name stack after the update.
-
-  Returns:
-    popped: A list of names popped from the name stack as part of the update.
-    pushed: A list of names pushed to the name stack as part of the update.
-  """
-  common_prefix_idx = 0
-  for i, (old, new) in enumerate(unsafe_zip(old_name_stack, new_name_stack)):
-    if old == new:
-      common_prefix_idx = i+1
-    else:
-      break
-  return old_name_stack[common_prefix_idx:], new_name_stack[common_prefix_idx:]
-
-
 def jaxpr_subcomp(
     ctx: LoweringContext, jaxpr: jax_core.Jaxpr, *args: ir.Value
 ) -> Sequence[ir.Value]:
@@ -618,9 +595,6 @@ def jaxpr_subcomp(
     block_shape_env[invar] = bs
   map(write_env, jaxpr.invars, args)
 
-  current_name_stack: list[str] = []
-  # TODO(justinfu): Handle transform scopes.
-  current_name_stack.extend([scope.name for scope in ctx.name_stack.stack])
   for eqn in jaxpr.eqns:
     invals = map(read_env, eqn.invars)
     source_info = eqn.source_info.replace(
@@ -641,17 +615,6 @@ def jaxpr_subcomp(
             [v.aval for v in eqn.outvars],
             block_shapes,
         )
-
-        # Insert trace_start and trace_stop ops on named_scope boundaries.
-        name_stack = [scope.name for scope in source_info.name_stack.stack]
-        popped, pushed = _compute_name_stack_updates(
-            current_name_stack, name_stack)
-        current_name_stack = name_stack
-        for _ in popped:
-          _trace_stop_lowering_rule(rule_context)
-        for name in pushed:
-          _trace_start_lowering_rule(rule_context, message=name, level=10)
-
         try:
           ans = lowering_rules[eqn.primitive](
               rule_context, *invals, **eqn.params
diff --git a/jax/_src/pallas/pallas_call.py b/jax/_src/pallas/pallas_call.py
index 5adcf2ac7..e0332bae3 100644
--- a/jax/_src/pallas/pallas_call.py
+++ b/jax/_src/pallas/pallas_call.py
@@ -65,7 +65,7 @@ def _maybe_dynamic_slice(start_idx, block_shape, value, is_indexing):
     return value
   assert is_indexing is not None
   start_idx = tuple(jnp.asarray(s, dtype=jnp.int32) for s in start_idx)
-  output = _padded_dynamic_slice(value, start_idx, block_shape)
+  output = lax.dynamic_slice(value, start_idx, slice_sizes=block_shape)
   squeeze_dims = tuple(np.arange(len(is_indexing))[np.array(is_indexing,
                                                             dtype=np.bool_)])
   return lax.squeeze(output, squeeze_dims)
@@ -81,7 +81,7 @@ def _maybe_dynamic_update_slice(start_idx, block_shape, value, update,
                          if not b)
   update = lax.broadcast_in_dim(update, block_shape, broadcast_dims)
   assert update.shape == block_shape
-  return _padded_dynamic_update_slice(value, update, start_idx)
+  return lax.dynamic_update_slice(value, update, start_idx)
 
 def _uninitialized_value(shape, dtype):
   if jnp.issubdtype(dtype, jnp.floating):
@@ -90,48 +90,6 @@ def _uninitialized_value(shape, dtype):
     return jnp.full(shape, jnp.iinfo(dtype).min, dtype)
   raise NotImplementedError(dtype)
 
-def _padded_dynamic_slice(
-    value: jax.Array, start_idx: jax.Array, slice_sizes: tuple[int, ...]):
-  """Performs a dynamic_slice with padding if required."""
-  if not value.shape:  # Don't pad for scalars.
-    return lax.dynamic_slice(value, start_idx, slice_sizes)
-  requires_pad = jnp.array(
-      [i+b > v for i, b, v in zip(start_idx, slice_sizes, value.shape)])
-  requires_pad = jnp.any(requires_pad)
-  pad_value = _uninitialized_value(shape=(), dtype=value.dtype)
-  def padded_op():
-    pad_width = tuple((0, v+b) for v, b in zip(value.shape, slice_sizes))
-    padded_value = jnp.pad(value, pad_width, constant_values=pad_value)
-    return lax.dynamic_slice(padded_value, start_idx, slice_sizes)
-  result = lax.cond(requires_pad & (len(value.shape) > 0),
-                    padded_op,
-                    lambda: lax.dynamic_slice(value, start_idx, slice_sizes)
-           )
-  return result
-
-def _padded_dynamic_update_slice(
-    value: jax.Array, updates: jax.Array, start_idx: jax.Array):
-  """Performs a dynamic_update_slice with padding if required."""
-  if not value.shape:  # Don't pad for scalars.
-    return lax.dynamic_update_slice(value, updates, start_idx)
-  slice_sizes = updates.shape
-  requires_pad = jnp.array(
-      [i+b > v for i, b, v in zip(start_idx, slice_sizes, value.shape)])
-  requires_pad = jnp.any(requires_pad)
-  pad_value = _uninitialized_value(shape=(), dtype=value.dtype)
-  def padded_op():
-    pad_width = tuple((0, v+b) for v, b in zip(value.shape, slice_sizes))
-    padded_value = jnp.pad(value, pad_width, constant_values=pad_value)
-    padded_value = lax.dynamic_update_slice(padded_value, updates, start_idx)
-    return lax.slice(
-        padded_value, (0,) * len(value.shape), value.shape
-    )
-  result = lax.cond(requires_pad & (len(value.shape) > 0),
-                    padded_op,
-                    lambda: lax.dynamic_update_slice(value, updates, start_idx),
-           )
-  return result
-
 def _get_next_indices(grid, indices):
   next_indices = []
   carry = True
diff --git a/jax/_src/pallas/triton/BUILD b/jax/_src/pallas/triton/BUILD
index d8b4c6ed3..bd14339ae 100644
--- a/jax/_src/pallas/triton/BUILD
+++ b/jax/_src/pallas/triton/BUILD
@@ -35,29 +35,10 @@ py_library_providing_imports_info(
     deps = [
         ":lowering",
         ":pallas_call_registration",
-        ":primitives",
         "//jax/_src/lib",
     ],
 )
 
-pytype_strict_library(
-    name = "primitives",
-    srcs = ["primitives.py"],
-    deps = [
-        ":lowering",
-        "//jax",
-        "//jax:ad_util",
-        "//jax:api_util",
-        "//jax:core",
-        "//jax:mlir",
-        "//jax:partial_eval",
-        "//jax:source_info_util",
-        "//jax:util",
-        "//jax/_src/lib",
-        "//jax/_src/pallas",
-    ] + py_deps("numpy"),
-)
-
 pytype_strict_library(
     name = "lowering",
     srcs = ["lowering.py"],
diff --git a/jax/_src/pallas/triton/__init__.py b/jax/_src/pallas/triton/__init__.py
index d881f99a0..1cd25df19 100644
--- a/jax/_src/pallas/triton/__init__.py
+++ b/jax/_src/pallas/triton/__init__.py
@@ -15,7 +15,6 @@
 """Contains Triton-specific pallas modules."""
 
 from jax._src.lib import gpu_triton as triton_kernel_call_lib
-from jax._src.pallas.triton.primitives import approx_tanh
 
 
 try:
diff --git a/jax/_src/pallas/triton/lowering.py b/jax/_src/pallas/triton/lowering.py
index cdad407fd..d298b1a9f 100644
--- a/jax/_src/pallas/triton/lowering.py
+++ b/jax/_src/pallas/triton/lowering.py
@@ -21,7 +21,7 @@ import dataclasses
 import functools
 import math
 import operator
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable
 
 import jax
 from jax import lax
@@ -56,12 +56,11 @@ from jax._src.util import split_list
 import jax.numpy as jnp
 import numpy as np
 
+
 # TODO(sharadmv): Enable type checking.
 # mypy: ignore-errors
 # pytype: skip-file
 
-_T = TypeVar("_T")
-
 map, unsafe_map = util.safe_map, map
 zip, unsafe_zip = util.safe_zip, zip
 
@@ -168,13 +167,6 @@ def _bcast(
 triton_lowering_rules = {}
 
 
-def register_lowering(primitive: jax_core.Primitive) -> Callable[[_T], _T]:
-  def wrapper(fn):
-    triton_lowering_rules[primitive] = fn
-    return fn
-  return wrapper
-
-
 def _process_grid_to_3d_grid(grid_mapping: GridMapping):
   launch_grid = []
   launch_grid_to_pallas_grid = []
diff --git a/jax/_src/pallas/triton/primitives.py b/jax/_src/pallas/triton/primitives.py
deleted file mode 100644
index da2eaee43..000000000
--- a/jax/_src/pallas/triton/primitives.py
+++ /dev/null
@@ -1,68 +0,0 @@
-# Copyright 2024 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Module for GPU-specific JAX primitives."""
-
-from __future__ import annotations
-
-import jax
-from jax import core as jax_core
-from jax._src.lib.triton import dialect as tt_dialect
-from jax._src.pallas.triton import lowering
-import jax.numpy as jnp
-
-
-def approx_tanh(x: jax.typing.ArrayLike) -> jax.Array:
-  r"""Elementwise approximate hyperbolic tangent: :math:`\mathrm{tanh}(x)`.
-
-  See
-  https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#floating-point-instructions-tanh.
-  """
-  return approx_tanh_p.bind(x)
-
-
-approx_tanh_p = jax_core.Primitive("approx_tanh_p")
-
-
-@approx_tanh_p.def_abstract_eval
-def _approx_tanh_abstract_eval(
-    x_aval: jax_core.ShapedArray,
-) -> jax_core.ShapedArray:
-  if jnp.dtype(x_aval.dtype) not in (jnp.float16, jnp.bfloat16, jnp.float32):
-    raise TypeError(f"approx_tanh does not accept {x_aval.dtype} arrays")
-  return x_aval
-
-
-@lowering.register_lowering(approx_tanh_p)
-def _approx_tanh_lowering(ctx: lowering.LoweringContext, x):
-  [x_aval] = ctx.avals_in
-  if x_aval.dtype == jnp.float16:
-    asm = "tanh.approx.f16 $0, $1;"
-    constraint = "h"
-  elif x_aval.dtype == jnp.bfloat16:
-    asm = "tanh.approx.bf16 $0, $1;"
-    constraint = "h"
-  elif x_aval.dtype == jnp.float32:
-    asm = "tanh.approx.f32 $0, $1;"
-    constraint = "f"
-  else:
-    raise NotImplementedError(f"Unsupported dtype: {x_aval.dtype}")
-  return tt_dialect.elementwise_inline_asm(
-      [x.type],
-      asm,
-      constraints=f"={constraint},{constraint}",
-      pure=True,
-      packed_element=1,
-      args=[x],
-  )
diff --git a/jax/_src/tpu_custom_call.py b/jax/_src/tpu_custom_call.py
index 3699e7ee4..bcf7957ee 100644
--- a/jax/_src/tpu_custom_call.py
+++ b/jax/_src/tpu_custom_call.py
@@ -26,6 +26,7 @@ import os
 import time
 from typing import Any, Callable
 
+from absl import flags
 import jax
 from jax import core
 from jax._src import config
@@ -41,11 +42,7 @@ from jaxlib.mlir.dialects import stablehlo
 from jaxlib.mlir.passmanager import PassManager
 import numpy as np
 
-try:
-  from absl import flags
-  FLAGS = flags.FLAGS
-except ImportError:
-  FLAGS = {}
+FLAGS = flags.FLAGS
 
 _MOSAIC_USE_PYTHON_PIPELINE = config.define_bool_state(
     name="mosaic_use_python_pipeline",
@@ -324,12 +321,7 @@ def _lower_tpu_kernel(
     pipeline.run(module.operation)
     dump_mlir(module, "post-simplify")
 
-    try:
-      on_device_checks = FLAGS["xla_mosaic_on_device_checks"].value
-    except KeyError:
-      on_device_checks = False
-
-    if checks := on_device_checks:
+    if checks := FLAGS["xla_mosaic_on_device_checks"].value:
       checks = set(checks.split(","))
       if checks == {"bounds"}:  # We only support one kind of checks now.
         pipeline = PassManager.parse(
diff --git a/jax/experimental/array_api/__init__.py b/jax/experimental/array_api/__init__.py
index 7050a4e3c..d8c42bc56 100644
--- a/jax/experimental/array_api/__init__.py
+++ b/jax/experimental/array_api/__init__.py
@@ -180,6 +180,7 @@ from jax.experimental.array_api._manipulation_functions import (
     roll as roll,
     squeeze as squeeze,
     stack as stack,
+    tile as tile,
     unstack as unstack,
 )
 
diff --git a/jax/experimental/array_api/_manipulation_functions.py b/jax/experimental/array_api/_manipulation_functions.py
index 3a1845909..e0e7a8bcc 100644
--- a/jax/experimental/array_api/_manipulation_functions.py
+++ b/jax/experimental/array_api/_manipulation_functions.py
@@ -83,6 +83,12 @@ def stack(arrays: tuple[Array, ...] | list[Array], /, *, axis: int = 0) -> Array
   dtype = _result_type(*arrays)
   return jax.numpy.stack(arrays, axis=axis, dtype=dtype)
 
+
+def tile(x: Array, repetitions: tuple[int], /) -> Array:
+  """Constructs an array by tiling an input array."""
+  return jax.numpy.tile(x, repetitions)
+
+
 def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]:
   """Splits an array in a sequence of arrays along the given axis."""
   return jax.numpy.unstack(x, axis=axis)
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 35272f250..5540219ec 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -414,9 +414,6 @@ def export(fun_jax: Callable,
 
       symbolic_scope: tuple[_shape_poly.SymbolicScope, tree_util.KeyPath] | None = None
       for k_path, aval in tree_util.tree_flatten_with_path((args_specs, kwargs_specs))[0]:
-        # Static args may has no `shape` attribute.
-        if not hasattr(aval, "shape"):
-          continue
         for d in aval.shape:
           if _shape_poly.is_symbolic_dim(d):
             if symbolic_scope is None:
@@ -793,7 +790,7 @@ def _check_lowering(lowering) -> None:
 # Their backwards compatibility is tested by back_compat_test.py.
 _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE = {
     "Sharding", "SPMDFullToShardShape", "SPMDShardToFullShape",
-    "dynamic_ducc_fft", "cu_threefry2x32",
+    "ducc_fft", "dynamic_ducc_fft", "cu_threefry2x32",
     # cholesky on CPU
     "lapack_spotrf", "lapack_dpotrf", "lapack_cpotrf", "lapack_zpotrf",
     # eigh on CPU
@@ -812,7 +809,7 @@ _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE = {
     "lapack_sgesdd", "lapack_dgesdd", "lapack_cgesdd", "lapack_zgesdd",
     # qr on GPU
     "cusolver_geqrf", "cublas_geqrf_batched",
-    "cusolver_orgqr",
+    "cusolver_geqrf", "cusolver_orgqr",
     # qr and svd on TPU
     "Qr", "ProductOfElementaryHouseholderReflectors",
     # triangular_solve on CPU
diff --git a/jax/experimental/mosaic/gpu/__init__.py b/jax/experimental/mosaic/gpu/__init__.py
index 9e47f2b18..d1b3c935e 100644
--- a/jax/experimental/mosaic/gpu/__init__.py
+++ b/jax/experimental/mosaic/gpu/__init__.py
@@ -1,4 +1,3 @@
-from collections.abc import Callable
 # Copyright 2024 The JAX Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,7 +25,6 @@ from typing import Any, Sequence
 
 import jax
 from jax._src import config
-from jax._src import core as jax_core
 from jax._src.interpreters import mlir
 from jax._src.lib import xla_client
 from jax._src.lib import mosaic_gpu as mosaic_gpu_lib
@@ -60,9 +58,6 @@ else:
 PTXAS_PATH = os.path.join(CUDA_ROOT, "bin/ptxas")
 NVDISASM_PATH = os.path.join(CUDA_ROOT, "bin/nvdisasm")
 
-TMA_DESCRIPTOR_BYTES = 128
-TMA_DESCRIPTOR_ALIGNMENT = 64
-
 
 c = mgpu.c  # This is too common to fully qualify.
 
@@ -102,13 +97,11 @@ mosaic_gpu_p.multiple_results = True
 
 
 @mosaic_gpu_p.def_abstract_eval
-def _mosaic_gpu_abstract_eval(*_, module, out_types, gmem_scratch_bytes):
-  del module, gmem_scratch_bytes  # Unused.
+def _mosaic_gpu_abstract_eval(*_, module, out_types):
   return [jax._src.core.ShapedArray(t.shape, t.dtype) for t in out_types]
 
 
-def _mosaic_gpu_lowering_rule(ctx, *args, module, out_types, gmem_scratch_bytes):
-  del out_types  # Unused.
+def _mosaic_gpu_lowering_rule(ctx, *args, module, out_types):
   runtime_path = (
       pathlib.Path(mosaic_gpu_lib._mosaic_gpu_ext.__file__).parent.parent.parent
       / "mosaic" / "gpu" / "libmlir_cuda_runtime.so"
@@ -134,16 +127,11 @@ def _mosaic_gpu_lowering_rule(ctx, *args, module, out_types, gmem_scratch_bytes)
   )  # pytype: disable=attribute-error
   op = mlir.custom_call(
       "mosaic_gpu",
-      result_types=[
-          *(mlir.aval_to_ir_type(aval) for aval in ctx.avals_out),
-          mlir.aval_to_ir_type(
-              jax_core.ShapedArray((gmem_scratch_bytes,), np.uint8)
-          ),
-      ],
+      result_types=[mlir.aval_to_ir_type(aval) for aval in ctx.avals_out],
       operands=args,
       backend_config=ptr_bytes,
   )
-  return op.results[:-1]  # Skip the scratch space.
+  return op.results
 
 mlir.register_lowering(mosaic_gpu_p, _mosaic_gpu_lowering_rule, "cuda")
 
@@ -239,12 +227,7 @@ OnDeviceProfiler = profiler.OnDeviceProfiler
 @dataclasses.dataclass()
 class LaunchContext:
   launch_op: gpu.LaunchOp
-  gmem_scratch_ptr: ir.Value
   profiler: OnDeviceProfiler | None = None
-  next_scratch_offset: int = 0
-  host_scratch_init: list[Callable[[ir.Value], None]] = dataclasses.field(
-      default_factory=list, init=False
-  )
   tma_descriptors: dict[
       tuple[ir.Value, tuple[int, ...], int | None, tuple[MemRefTransform, ...]],
       ir.Value,
@@ -258,37 +241,6 @@ class LaunchContext:
     else:
       yield
 
-  def _alloc_scratch(
-      self,
-      size: int,
-      alignment: int | None = None,
-      host_init: Callable[[ir.Value], None] = lambda _: None,
-      device_init: Callable[[ir.Value], Any] = lambda x: x,
-  ) -> ir.Value:
-    """Allocates a GMEM scratch buffer.
-
-    The buffer is initialized on the host and then copied to GMEM before the
-    kernel launch.
-    """
-    i8 = ir.IntegerType.get_signless(8)
-    ptr_ty = ir.Type.parse("!llvm.ptr")
-    if alignment is None:
-      alignment = size
-    if self.next_scratch_offset % alignment:
-      raise NotImplementedError  # TODO(apaszke): Pad to match alignment
-    alloc_base = self.next_scratch_offset
-    self.next_scratch_offset += size
-    def host_init_wrapped(host_ptr):
-      with ir.InsertionPoint(self.launch_op):
-        host_init(
-            llvm.getelementptr(ptr_ty, host_ptr, [], [alloc_base], i8)
-        )
-    self.host_scratch_init.append(host_init_wrapped)
-    with ir.InsertionPoint.at_block_begin(self.launch_op.body.blocks[0]):
-      return device_init(llvm.getelementptr(
-          ptr_ty, self.gmem_scratch_ptr, [], [alloc_base], i8
-      ))
-
   def _get_tma_desc(
       self,
       ref,
@@ -313,42 +265,13 @@ class LaunchContext:
       with ir.InsertionPoint(self.launch_op):
         for t in gmem_transform:
           ref = t.apply(ref)
-        ref_ty = ir.MemRefType(ref.type)
-
-        i64 = ir.IntegerType.get_signless(64)
-        ptr_ty = ir.Type.parse("!llvm.ptr")
-        def init_tma_desc(host_ptr):
-          _, offset, *sizes_and_strides = memref.extract_strided_metadata(ref)
-          aligned_ptr_idx = memref.extract_aligned_pointer_as_index(ref)
-          as_i64 = lambda i: arith.index_cast(i64, i)
-          alloc_ptr = llvm.inttoptr(ptr_ty, as_i64(aligned_ptr_idx))
-          llvm_dyn = -2147483648  # TODO(apaszke): Improve the MLIR bindings...
-          base_ptr = llvm.getelementptr(
-              ptr_ty, alloc_ptr, [as_i64(offset)], [llvm_dyn], ref_ty.element_type,
-          )
-          rank = ref_ty.rank
-          assert rank * 2 == len(sizes_and_strides)
-          args = [
-              host_ptr,
-              base_ptr,
-              c(utils.bytewidth(ref_ty.element_type), i64),
-              c(rank, i64),
-              utils.pack_array([as_i64(i) for i in sizes_and_strides[:rank]]),
-              utils.pack_array([as_i64(i) for i in sizes_and_strides[rank:]]),
-              c(0 if swizzle is None else swizzle, i64),
-              utils.pack_array([c(v, i64) for v in transformed_slice_shape]),
-          ]
-          func.call([], "mosaic_gpu_init_tma_desc", args)
-        def cast_tma_desc(device_ptr):
-          nvvm.prefetch_tensormap(device_ptr)
-          return builtin.unrealized_conversion_cast(
-              [tensor_map_ty], [device_ptr]
-          )
-        tma_desc = self._alloc_scratch(
-            TMA_DESCRIPTOR_BYTES,
-            alignment=TMA_DESCRIPTOR_ALIGNMENT,
-            host_init=init_tma_desc,
-            device_init=cast_tma_desc,
+        ref_unranked = memref.cast(
+            ir.UnrankedMemRefType.get(ref_ty.element_type, None), ref
+        )
+        tma_desc = nvgpu.tma_create_descriptor(
+            tensor_map_ty,
+            ref_unranked,
+            [c(s, index) for s in transformed_slice_shape],
         )
         self.tma_descriptors[tma_desc_key] = tma_desc
     return tma_desc
@@ -455,14 +378,18 @@ class LaunchContext:
     nvvm.cp_async_bulk_wait_group(allow_groups, read=await_read_only)
     gpu.barrier()  # Groups are supposedly tracked per-thread
 
+  def _prefetch_tma_descs(self):
+    with ir.InsertionPoint(self.launch_op.body.blocks[0]):
+      with mgpu.once():
+        for desc in self.tma_descriptors.values():
+          nvgpu.tma_prefetch_descriptor(desc)
+
 
-# TODO(apaszke): Inline this
 @contextlib.contextmanager
 def _launch(
     token,
     grid,
     block,
-    gmem_scratch_ptr,
     smem_buffers,
     profiler_spec: profiler.ProfilerSpec | None = None,
     maybe_prof_buffer: ir.Value | None = None,
@@ -522,7 +449,7 @@ def _launch(
     else:
       prof = None
     smem_ref_tree = jax.tree.unflatten(smem_buffer_tree, smem_refs)
-    yield LaunchContext(launch_op, gmem_scratch_ptr, prof), smem_ref_tree
+    yield LaunchContext(launch_op, prof), smem_ref_tree
     if prof is not None:
       prof.finalize(grid=grid)
     gpu.terminator()
@@ -539,8 +466,6 @@ def as_gpu_kernel(
 ):
   ptr_ty = ir.Type.parse("!llvm.ptr")
   token_ty = ir.Type.parse("!gpu.async.token")
-  i8 = ir.IntegerType.get_signless(8)
-  i64 = ir.IntegerType.get_signless(64)
 
   def _shape_to_ref_ty(shape: jax.ShapeDtypeStruct) -> ir.MemRefType:
     return ir.MemRefType.get(shape.shape, mlir.dtype_to_ir_type(shape.dtype))
@@ -564,46 +489,20 @@ def as_gpu_kernel(
 
   module = ir.Module.create()
   with ir.InsertionPoint(module.body):
-    _declare_runtime_functions()
-    gmem_scratch_bytes = 0
     @func.FuncOp.from_py_func(ptr_ty, ptr_ty)
     def main(token_ptr, buffers):
-      nonlocal gmem_scratch_bytes
       token = builtin.unrealized_conversion_cast([token_ty], [token_ptr])
       arg_refs = []
-      i = -1
       for i, ref_ty in enumerate([*in_ref_tys, *out_ref_tys]):
         ptr = llvm.LoadOp(ptr_ty, llvm.GEPOp(ptr_ty, buffers, [], [i], ptr_ty))
         arg_refs.append(utils.ptr_as_memref(ptr, ir.MemRefType(ref_ty)))
-      gmem_scratch_ptr = llvm.LoadOp(
-          ptr_ty, llvm.GEPOp(ptr_ty, buffers, [], [i + 1], ptr_ty)
-      )
       in_refs = arg_refs[:len(in_ref_tys)]
       out_refs = arg_refs[len(in_ref_tys):]
       prof_buffer = out_refs.pop() if prof_spec is not None else None
       with _launch(
-          token, grid, block, gmem_scratch_ptr, smem_scratch_shape,
-          prof_spec, prof_buffer
+          token, grid, block, smem_scratch_shape, prof_spec, prof_buffer
       ) as (launch_ctx, smem_refs):
         body(launch_ctx, *in_refs, *out_refs, smem_refs)
-        gmem_scratch_bytes = launch_ctx.next_scratch_offset
-      # Allocate and initialize the host buffer right before the launch.
-      # Note that we couldn't do that before, because we had to run the body
-      # to learn what the scratch contains.
-      with ir.InsertionPoint(launch_ctx.launch_op):
-        host_scratch_ptr = llvm.alloca(ptr_ty, c(gmem_scratch_bytes, i64), i8)
-        for init_callback in launch_ctx.host_scratch_init:
-          init_callback(host_scratch_ptr)
-        func.call(
-            [],
-            "mosaic_gpu_memcpy_async_h2d",
-            [
-                gmem_scratch_ptr,
-                host_scratch_ptr,
-                c(gmem_scratch_bytes, i64),
-                token_ptr,
-            ],
-        )
     main.func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
   module.operation.verify()
 
@@ -624,12 +523,7 @@ def as_gpu_kernel(
   pass_manager.run(module.operation)
 
   def bind(*args):
-    return mosaic_gpu_p.bind(
-        *args,
-        out_types=out_shape,
-        module=module,
-        gmem_scratch_bytes=gmem_scratch_bytes,
-    )
+    return mosaic_gpu_p.bind(*args, out_types=out_shape, module=module)
 
   if prof_spec is not None:
     @jax.jit
@@ -658,21 +552,6 @@ def as_gpu_kernel(
     return kernel
 
 
-def _declare_runtime_functions():
-  """Declares the runtime functions that can be used by the generated code."""
-  ptr_ty = ir.Type.parse("!llvm.ptr")
-  i64 = ir.IntegerType.get_signless(64)
-  arg_tys = [ptr_ty, ptr_ty, i64, i64, ptr_ty, ptr_ty, i64, ptr_ty]
-  init_tma_desc_type = ir.FunctionType.get(arg_tys, [])
-  func.FuncOp(
-      "mosaic_gpu_init_tma_desc", init_tma_desc_type, visibility="private"
-  )
-  memcpy_async_type = ir.FunctionType.get([ptr_ty, ptr_ty, i64, ptr_ty], [])
-  func.FuncOp(
-      "mosaic_gpu_memcpy_async_h2d", memcpy_async_type, visibility="private"
-  )
-
-
 def dump_low_level(module):
   dump_ptx = mosaic_gpu_dump_ptx.value
   dump_ptxas = mosaic_gpu_dump_ptxas.value
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index 9e704533a..103187e23 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -69,19 +69,6 @@ def ptr_as_memref(ptr, memref_ty: ir.MemRefType):
   return builtin.unrealized_conversion_cast([memref_ty], [desc])
 
 
-def pack_array(values):
-  if not values:
-    raise ValueError("Empty array")
-  elem_ty = values[0].type
-  i64 = ir.IntegerType.get_signless(64)
-  ptr_ty = ir.Type.parse("!llvm.ptr")
-  arr_ptr = llvm.alloca(ptr_ty, c(len(values), i64), elem_ty)
-  for i, v in enumerate(values):
-    elem_ptr = llvm.getelementptr(ptr_ty, arr_ptr, [], [i], elem_ty)
-    llvm.store(v, elem_ptr)
-  return arr_ptr
-
-
 def get_contiguous_strides(xs):
   strides_ret = []
   stride = 1
diff --git a/jax/experimental/pallas/gpu.py b/jax/experimental/pallas/gpu.py
index fef855613..e58bd498b 100644
--- a/jax/experimental/pallas/gpu.py
+++ b/jax/experimental/pallas/gpu.py
@@ -14,6 +14,5 @@
 
 """Contains Triton specific Pallas functions."""
 from jax._src.pallas import triton
-from jax._src.pallas.triton import approx_tanh
 get_compute_capability = triton.get_compute_capability
 del triton
diff --git a/jaxlib/cpu/cpu_kernels.cc b/jaxlib/cpu/cpu_kernels.cc
index b7551cc61..dc5d657a9 100644
--- a/jaxlib/cpu/cpu_kernels.cc
+++ b/jaxlib/cpu/cpu_kernels.cc
@@ -15,7 +15,7 @@ limitations under the License.
 
 // This file is not used by JAX itself, but exists to assist with running
 // JAX-generated HLO code from outside of JAX.
-#include <complex>
+
 #include "jaxlib/cpu/ducc_fft_kernels.h"
 #include "jaxlib/cpu/lapack_kernels.h"
 #include "xla/service/custom_call_target_registry.h"
@@ -105,6 +105,8 @@ XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(
     "lapack_cgees", ComplexGees<std::complex<float>>::Kernel, "Host");
 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(
     "lapack_zgees", ComplexGees<std::complex<double>>::Kernel, "Host");
+XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(
+    "ducc_fft", DuccFft, "Host");
 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(
     "dynamic_ducc_fft", DynamicDuccFft, "Host");
 
diff --git a/jaxlib/cpu/ducc_fft.cc b/jaxlib/cpu/ducc_fft.cc
index a8f0490ac..33e73c5f4 100644
--- a/jaxlib/cpu/ducc_fft.cc
+++ b/jaxlib/cpu/ducc_fft.cc
@@ -48,8 +48,9 @@ nb::bytes BuildDynamicDuccFftDescriptor(
 
 nb::dict Registrations() {
   nb::dict dict;
-  // TODO(b/311175955): this must be kept until May 2024 for backwards
+  // TODO(b/287702203): this must be kept until EOY 2023 for backwards
   // of serialized functions using fft.
+  dict["ducc_fft"] = EncapsulateFunction(DuccFft);
   dict["dynamic_ducc_fft"] = EncapsulateFunction(DynamicDuccFft);
   return dict;
 }
diff --git a/jaxlib/cpu/ducc_fft.fbs b/jaxlib/cpu/ducc_fft.fbs
index bc8572ad5..a58e1dc7c 100644
--- a/jaxlib/cpu/ducc_fft.fbs
+++ b/jaxlib/cpu/ducc_fft.fbs
@@ -26,6 +26,17 @@ enum DuccFftType : byte {
   R2C = 2,
 }
 
+table DuccFftDescriptor {
+  dtype:DuccFftDtype;
+  fft_type:DuccFftType;
+  shape:[uint64];
+  strides_in:[uint64];
+  strides_out:[uint64];
+  axes:[uint32];
+  forward:bool;
+  scale:double;
+}
+
 table DynamicDuccFftDescriptor {
   ndims:uint32;
   dtype:DuccFftDtype;
@@ -34,4 +45,4 @@ table DynamicDuccFftDescriptor {
   forward:bool;
 }
 
-root_type DynamicDuccFftDescriptor;
+root_type DuccFftDescriptor;
diff --git a/jaxlib/cpu/ducc_fft_kernels.cc b/jaxlib/cpu/ducc_fft_kernels.cc
index 12f8327b6..eab35905f 100644
--- a/jaxlib/cpu/ducc_fft_kernels.cc
+++ b/jaxlib/cpu/ducc_fft_kernels.cc
@@ -98,8 +98,23 @@ void DuccFftImpl(void *out, void *operand, jax::DuccFftDtype dtype,
 }  // namespace
 
 
-// TODO(b/311175955): this must be kept until May 2024 for backwards
+// TODO(b/287702203): this must be kept until EOY 2023 for backwards
 // of serialized functions using fft.
+void DuccFft(void *out, void **in, XlaCustomCallStatus *) {
+  const DuccFftDescriptor *descriptor = GetDuccFftDescriptor(in[0]);
+  shape_t shape(descriptor->shape()->begin(), descriptor->shape()->end());
+  stride_t strides_in(descriptor->strides_in()->begin(),
+                      descriptor->strides_in()->end());
+  stride_t strides_out(descriptor->strides_out()->begin(),
+                       descriptor->strides_out()->end());
+  shape_t axes(descriptor->axes()->begin(), descriptor->axes()->end());
+
+  DuccFftImpl(out, in[1], descriptor->dtype(), descriptor->fft_type(),
+              shape, strides_in, strides_out, axes,
+              descriptor->forward(), descriptor->scale());
+}
+
+
 void DynamicDuccFft(void *out, void **in, XlaCustomCallStatus *) {
   // in[0]=descriptor, in[1]=operand,
   // in[2]=shape, in[3]=strides_in, in[4]=strides_out, in[5]=scale.
diff --git a/jaxlib/cpu/ducc_fft_kernels.h b/jaxlib/cpu/ducc_fft_kernels.h
index 13d0b1d40..a3bf6cf46 100644
--- a/jaxlib/cpu/ducc_fft_kernels.h
+++ b/jaxlib/cpu/ducc_fft_kernels.h
@@ -20,9 +20,10 @@ limitations under the License.
 
 namespace jax {
 
-
-// TODO(b/311175955): this must be kept until May 2024 for backwards
+// TODO(b/287702203): this must be kept until EOY 2023 for backwards
 // of serialized functions using fft.
+void DuccFft(void* out, void** in, XlaCustomCallStatus*);
+
 void DynamicDuccFft(void* out, void** in, XlaCustomCallStatus*);
 
 }  // namespace jax
diff --git a/jaxlib/cuda/BUILD b/jaxlib/cuda/BUILD
index 1b976ab32..99449a5d9 100644
--- a/jaxlib/cuda/BUILD
+++ b/jaxlib/cuda/BUILD
@@ -37,7 +37,6 @@ cc_library(
     defines = ["JAX_GPU_CUDA=1"],
     visibility = ["//visibility:public"],
     deps = [
-        "@xla//xla/tsl/cuda:cupti",
         "@local_config_cuda//cuda:cuda_headers",
         "@local_config_cuda//cuda:cudnn_header",
     ],
diff --git a/jaxlib/mlir/_mlir_libs/BUILD.bazel b/jaxlib/mlir/_mlir_libs/BUILD.bazel
index e0d499439..d947ef6ee 100644
--- a/jaxlib/mlir/_mlir_libs/BUILD.bazel
+++ b/jaxlib/mlir/_mlir_libs/BUILD.bazel
@@ -219,12 +219,6 @@ pybind_extension(
         "-fexceptions",
         "-fno-strict-aliasing",
     ],
-    linkopts = select({
-        "@xla//xla/python:use_jax_cuda_pip_rpaths": [
-            "-Wl,-rpath,$$ORIGIN/../../../nvidia/cuda_runtime/lib",
-        ],
-        "//conditions:default": [],
-    }),
     visibility = ["//third_party/py/jax:__subpackages__"],
     deps = [
         ":jaxlib_mlir_capi_shared_library",
@@ -233,7 +227,6 @@ pybind_extension(
         "//jaxlib/mosaic/gpu:mlir_capi",
         "@nanobind",
         "@xla//xla/service:custom_call_status",
-        "@xla//xla/tsl/cuda:cudart",
     ],
 )
 
diff --git a/jaxlib/mosaic/gpu/BUILD b/jaxlib/mosaic/gpu/BUILD
index ff40b3165..d0615dd8a 100644
--- a/jaxlib/mosaic/gpu/BUILD
+++ b/jaxlib/mosaic/gpu/BUILD
@@ -87,20 +87,9 @@ cc_library(
     alwayslink = True,
 )
 
-cc_library(
-    name = "runtime",
-    srcs = ["runtime.cc"],
-    deps = [
-        "@local_config_cuda//cuda:cuda_headers",
-    ],
-)
-
 cc_binary(
     name = "libmlir_cuda_runtime.so",
-    srcs = [
-        "runtime.cc",
-        "@llvm-project//mlir:lib/ExecutionEngine/CudaRuntimeWrappers.cpp",
-    ],
+    srcs = ["@llvm-project//mlir:lib/ExecutionEngine/CudaRuntimeWrappers.cpp"],
     copts = ["-fvisibility=default"],
     linkopts = select({
         "@xla//xla/python:use_jax_cuda_pip_rpaths": [
diff --git a/jaxlib/mosaic/gpu/runtime.cc b/jaxlib/mosaic/gpu/runtime.cc
deleted file mode 100644
index d50ec8e63..000000000
--- a/jaxlib/mosaic/gpu/runtime.cc
+++ /dev/null
@@ -1,95 +0,0 @@
-/* Copyright 2024 The JAX Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include <cstdint>
-#include <cstdio>
-#include <cstdlib>
-#include "third_party/gpus/cuda/include/cuda.h"
-
-extern "C" {
-
-void mosaic_gpu_init_tma_desc(CUtensorMap *tma_desc, void *base_addr,
-                              int64_t elem_bytewidth, int64_t rank,
-                              int64_t *sizes, int64_t *strides,
-                              int64_t swizzle_bytes, int64_t *window_shape) {
-  CUtensorMapDataType data_type;
-  if (elem_bytewidth == 1) {
-    data_type = CU_TENSOR_MAP_DATA_TYPE_UINT8;
-  } else if (elem_bytewidth == 2) {
-    data_type = CU_TENSOR_MAP_DATA_TYPE_UINT16;
-  } else if (elem_bytewidth == 4) {
-    data_type = CU_TENSOR_MAP_DATA_TYPE_UINT32;
-  } else if (elem_bytewidth == 8) {
-    data_type = CU_TENSOR_MAP_DATA_TYPE_UINT64;
-  } else {
-    fprintf(stderr, "Unsupported element size: %ld\n", elem_bytewidth);
-    abort();
-  }
-  cuuint64_t tma_sizes[5] = {1, 1, 1, 1, 1};
-  for (int i = 0; i < rank; ++i) {
-    tma_sizes[i] = static_cast<cuuint64_t>(sizes[rank - i - 1]);
-  }
-  cuuint64_t tma_strides[5] = {1, 1, 1, 1, 1};
-  if (strides[rank - 1] != 1) {
-    fprintf(stderr, "Minormost stride must be 1, but got %ld\n",
-            strides[rank - 1]);
-    abort();
-  }
-  for (int i = 0; i < rank - 1; ++i) {  // We skip the implicit minor stride.
-    tma_strides[i] =
-        static_cast<cuuint64_t>(strides[rank - i - 2] * elem_bytewidth);
-  }
-  cuuint32_t tma_window_shape[5] = {1, 1, 1, 1, 1};
-  for (int64_t i = 0; i < rank; ++i) {
-    tma_window_shape[i] = static_cast<cuuint32_t>(window_shape[rank - i - 1]);
-  }
-  cuuint32_t element_strides[5] = {1, 1, 1, 1, 1};
-  CUtensorMapSwizzle swizzle;
-  if (swizzle_bytes == 0) {
-    swizzle = CU_TENSOR_MAP_SWIZZLE_NONE;
-  } else if (swizzle_bytes == 32) {
-    swizzle = CU_TENSOR_MAP_SWIZZLE_32B;
-  } else if (swizzle_bytes == 64) {
-    swizzle = CU_TENSOR_MAP_SWIZZLE_64B;
-  } else if (swizzle_bytes == 128) {
-    swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
-  } else {
-    fprintf(stderr, "Unsupported swizzle: %ld\n", swizzle_bytes);
-    abort();
-  }
-  CUresult result = cuTensorMapEncodeTiled(
-      tma_desc, data_type, rank, base_addr, tma_sizes, tma_strides,
-      tma_window_shape, element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle,
-      CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
-  if (result != CUDA_SUCCESS) {
-    const char *ptr = nullptr;
-    cuGetErrorString(result, &ptr);
-    fprintf(stderr, "cuTensorMapEncodeTiled failed: %s\n", ptr);
-    abort();
-  }
-}
-
-void mosaic_gpu_memcpy_async_h2d(CUdeviceptr dst, void *src, uint64_t bytes,
-                                 CUstream stream) {
-  CUresult result = cuMemcpyHtoDAsync(dst, src, bytes, stream);
-  if (result != CUDA_SUCCESS) {
-    const char *ptr = nullptr;
-    cuGetErrorString(result, &ptr);
-    fprintf(stderr, "cuMemcpyAsync failed: %s\n", ptr);
-    abort();
-  }
-}
-
-}
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index f4f4d2a46..04c1d8f7a 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -20,8 +20,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "822c341cdc1fd4f6a4cca47c11576f119acd868b"
-XLA_SHA256 = "e9d52805aea68d6d9b0b1c4dddf3805a3ceba004c2a0047eed8e1a18a0ac91aa"
+XLA_COMMIT = "e7c3bc72d2663edd0dc747b0e8ce623f5e840f04"
+XLA_SHA256 = "aa39a43662f66a55853ed450334de550f710a85f0f181fff19e479101bef6c12"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/jax/_src/internal_test_util/export_back_compat_test_data/cpu_ducc_fft.py b/jax/_src/internal_test_util/export_back_compat_test_data/cpu_ducc_fft.py
index 6418f3742..d92a90b36 100644
--- a/jax/_src/internal_test_util/export_back_compat_test_data/cpu_ducc_fft.py
+++ b/jax/_src/internal_test_util/export_back_compat_test_data/cpu_ducc_fft.py
@@ -15,6 +15,36 @@
 import datetime
 from numpy import array, float32, complex64
 
+# Pasted from the test output (see back_compat_test.py module docstring)
+data_2023_03_17 = dict(
+    testdata_version=1,
+    platform='cpu',
+    custom_call_targets=['ducc_fft'],
+    serialized_date=datetime.date(2023, 3, 17),
+    inputs=(array([[ 0.,  1.,  2.,  3.],
+       [ 4.,  5.,  6.,  7.],
+       [ 8.,  9., 10., 11.]], dtype=float32),),
+    expected_outputs=(array([[ 6.+0.j, -2.+2.j, -2.+0.j, -2.-2.j],
+       [22.+0.j, -2.+2.j, -2.+0.j, -2.-2.j],
+       [38.+0.j, -2.+2.j, -2.+0.j, -2.-2.j]], dtype=complex64),),
+    mlir_module_text=r"""
+module @jit_func {
+  func.func public @main(%arg0: tensor<3x4xf32> {jax.arg_info = "x", mhlo.sharding = "{replicated}"}) -> (tensor<3x4xcomplex<f32>> {jax.result_info = ""}) {
+    %0 = call @fft(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xcomplex<f32>>
+    return %0 : tensor<3x4xcomplex<f32>>
+  }
+  func.func private @fft(%arg0: tensor<3x4xf32>) -> tensor<3x4xcomplex<f32>> {
+    %0 = stablehlo.convert %arg0 : (tensor<3x4xf32>) -> tensor<3x4xcomplex<f32>>
+    %1 = stablehlo.constant dense<"0x18000000140024000000000008000C001000140007001800140000000000000154000000380000001C00000010000000000000000000F03F0000000001000000010000000200000004000000000000000100000000000000000000000200000004000000000000000100000000000000000000000200000003000000000000000400000000000000"> : tensor<136xui8>
+    %2 = stablehlo.custom_call @ducc_fft(%1, %0) {api_version = 2 : i32, operand_layouts = [dense<0> : tensor<1xindex>, dense<[1, 0]> : tensor<2xindex>], result_layouts = [dense<[1, 0]> : tensor<2xindex>]} : (tensor<136xui8>, tensor<3x4xcomplex<f32>>) -> tensor<3x4xcomplex<f32>>
+    return %2 : tensor<3x4xcomplex<f32>>
+  }
+}
+""",
+    mlir_module_serialized=b'ML\xefR\x03MLIRxxx-trunk\x00\x01\x1d\x05\x01\x05\x01\x03\x05\x03\r\x07\t\x0b\r\x0f\x11\x03\x99s\x15\x01?\x07\x0b\x0f\x17\x0b\x0b\x0b\x0b\x0f\x13\x0b33\x0b\x0b\x0f\x0b\x13\x0b\x0bK\x0b\x0b\x0b\x0b\x0b\x0b\x0b\x0b\x13\x0b\x035\x0b\x0b\x0f\x0b\x0bO\x0f\x1b\x0b\x0b\x0b\x0b\x0f\x13\x0b\x0b\x0b\x0b\x0bf\x04\x0b\x0b\x0b\x13/\x0f\x03\x15\x17\x17\x07\x17\x07\x17\x0b\x07\x13\x13\x02\xca\x05\x1f\x05\x13\x1d\x1b\x07\x17\x1d^\x03\x01\x05\x15\x05\x17\x05\x19\x05\x1b\x1d\'\x07\x03\x03\x03\x15\x05\x1d\x03\x0b\tK\x0b?\rW\x03]\x0f_\x03\x0b\tC\x0b?\rC\x03E\x0fc\x05\x1f\x05!\x1d!\x07\x05#\x03\x03%e\x05%\x05\'\x03\x11+g-A/i1G3k5m7G9q\x05)\x05+\x05-\x05/\x051\x053\x055\x057\x03\x03=E\x059#\x0b\x1d;\x03\x03a\x1d=\x03\x01\x1f\x13!\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x03M\r\x05OQSU\x1d?\x1dA\x1dC\x1dE\x03\x03Y\r\x03[A\x1dG\x1dI\x1dK\r\x01\x1dM\x1f\x07"\x02\x18\x00\x00\x00\x14\x00$\x00\x00\x00\x00\x00\x08\x00\x0c\x00\x10\x00\x14\x00\x07\x00\x18\x00\x14\x00\x00\x00\x00\x00\x00\x01T\x00\x00\x008\x00\x00\x00\x1c\x00\x00\x00\x10\x00\x00\x00\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x0b\x05\x1dO\x05\x01\x03\x05oI\x1f\x11\x11\x00\x00\x00\x00\x00\x00\x00\x00\x03\x03I)\x05\r\x11\r)\x05\r\x11\x05\t)\x03B\x04\x0f\x13\x11\x03\x03\x03\x01\x03\x05!)\x03\x05\t)\x03\t\t\x04\x8f\x05\x01\x11\x01\x13\x07\x03\x01\t\x03\x11\x01\x17\x05\x03\x05\x0b\x03\x03\x01\r\x07\x05;\x03\x01\x03\x01\x05\x04\x01\x03\x03\x03\x11\x05\x19\x05\x03\t\x13\x03\x03\x01\x07\x06\x1f\x03\x01\x03\x01\t\x03\x11#\x03\x07\x0b\x07\x11)\x03\x01\x05\x05\x03\x05\x04\x05\x03\x07\x06\x03\x01\x05\x01\x00\xc2\x0eQ\x13\x11\x0f\x0b!\x1b\x1d\x05\x1b\t\x03\x0f\x1f/!!)#\x1f\x19\x91\r\xaf\x83\x82\x04\x13\x1f\x15\x1d\x15\x13\x11\x1f\x19\x17\x15\x11\x0f\x0b\x11builtin\x00vhlo\x00module\x00func_v1\x00return_v1\x00convert_v1\x00constant_v1\x00custom_call_v1\x00call_v1\x00sym_name\x00arg_attrs\x00function_type\x00res_attrs\x00sym_visibility\x00jit_func\x00jit(func)/jit(main)/pjit[in_shardings=(UnspecifiedValue,) out_shardings=(UnspecifiedValue,) resource_env=None donated_invars=(False,) name=fft in_positional_semantics=(<_PositionalSemantics.GLOBAL: 1>,) out_positional_semantics=_PositionalSemantics.GLOBAL keep_unused=False inline=False]\x00third_party/py/jax/experimental/jax2tf/tests/back_compat_test.py\x00jit(func)/jit(main)/jit(fft)/convert_element_type[new_dtype=complex64 weak_type=False]\x00value\x00jit(func)/jit(main)/jit(fft)/fft[fft_type=FftType.FFT fft_lengths=(4,)]\x00api_version\x00backend_config\x00call_target_name\x00called_computations\x00has_side_effect\x00operand_layouts\x00output_operand_aliases\x00result_layouts\x00callee\x00\x00fft\x00jax.arg_info\x00x\x00mhlo.sharding\x00{replicated}\x00jax.result_info\x00main\x00public\x00private\x00ducc_fft\x00',
+    xla_call_module_version=4,
+)  # End paste
+
 # Pasted from the test output (see back_compat_test.py module docstring)
 data_2023_06_14 = dict(
     testdata_version=1,
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index 5667c3459..5ca850e7b 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -161,6 +161,7 @@ MAIN_NAMESPACE = {
   'tan',
   'tanh',
   'tensordot',
+  'tile',
   'tril',
   'triu',
   'trunc',
diff --git a/tests/export_back_compat_test.py b/tests/export_back_compat_test.py
index d2feef2e1..8a4a8304f 100644
--- a/tests/export_back_compat_test.py
+++ b/tests/export_back_compat_test.py
@@ -100,7 +100,7 @@ class CompatTest(bctu.CompatTestBase):
     # Add here all the testdatas that should cover the targets guaranteed
     # stable
     covering_testdatas = [
-        cpu_ducc_fft.data_2023_06_14,
+        cpu_ducc_fft.data_2023_03_17, cpu_ducc_fft.data_2023_06_14,
         cpu_cholesky_lapack_potrf.data_2023_06_19,
         cpu_eig_lapack_geev.data_2023_06_19,
         cpu_eigh_lapack_syev.data_2023_03_17,
@@ -142,6 +142,11 @@ class CompatTest(bctu.CompatTestBase):
       return lax.fft(x, fft_type="fft", fft_lengths=(4,))
 
     # TODO(b/311175955): Remove this test and the corresponding custom calls.
+
+    # An old lowering, with ducc_fft.
+    data = self.load_testdata(cpu_ducc_fft.data_2023_03_17)
+    self.run_one_test(func, data, expect_current_custom_calls=[])
+
     # A newer lowering, with dynamic_ducc_fft.
     data = self.load_testdata(cpu_ducc_fft.data_2023_06_14)
     # FFT no longer lowers to a custom call.
diff --git a/tests/export_test.py b/tests/export_test.py
index 01858950b..784351487 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -213,32 +213,6 @@ class JaxExportTest(jtu.JaxTestCase):
     f1 = export.call_exported(exp_f)
     self.assertAllClose(f(x), f1(x))
 
-  def test_jit_static_arg(self):
-
-    with self.subTest("static_argnames"):
-
-      @functools.partial(jax.jit, static_argnames=["c"])
-      def f(x, *, c):
-        return c * jnp.sin(x)
-
-      x = np.arange(4, dtype=np.float32)
-      exp_f = get_exported(f)(x, c=0.1)
-
-      f1 = export.call_exported(exp_f)
-      self.assertAllClose(f(x, c=0.1), f1(x))
-
-    with self.subTest("static_argnums"):
-
-      @functools.partial(jax.jit, static_argnums=[1])
-      def g(x, c):
-        return c * jnp.sin(x)
-
-      x = np.arange(4, dtype=np.float32)
-      exp_g = get_exported(g)(x, 0.1)
-
-      g1 = export.call_exported(exp_g)
-      self.assertAllClose(g(x, 0.1), g1(x))
-
   def test_call_exported_lambda(self):
     # When we export a lambda, the exported.fun_name is not a valid MLIR function name
     f = lambda x: jnp.sin(x)
@@ -521,6 +495,7 @@ class JaxExportTest(jtu.JaxTestCase):
             r"and found for 'w' \(args\[1\]\) scope .*", re.DOTALL)):
       get_exported(f)(x_poly_spec, y_poly_spec)
 
+
   @jtu.parameterized_filterable(
     kwargs=[
       dict(v=v)
diff --git a/tests/mosaic/BUILD b/tests/mosaic/BUILD
index ca98090b1..6830d621b 100644
--- a/tests/mosaic/BUILD
+++ b/tests/mosaic/BUILD
@@ -71,10 +71,6 @@ jax_test(
     disable_backends = DISABLED_BACKENDS,
     disable_configs = DISABLED_CONFIGS,
     main = "//third_party/py/jax/experimental/mosaic/gpu/examples:matmul.py",
-    tags = [
-        "manual",
-        "notap",
-    ],
     deps = [
         "//jax:mosaic_gpu",
     ] + py_deps("absl/testing") + py_deps("numpy"),
diff --git a/tests/pallas/pallas_call_tpu_test.py b/tests/pallas/pallas_call_tpu_test.py
index b70912102..d30878634 100644
--- a/tests/pallas/pallas_call_tpu_test.py
+++ b/tests/pallas/pallas_call_tpu_test.py
@@ -54,78 +54,6 @@ class PallasTPUTest(jtu.JaxTestCase):
     return pl.pallas_call(*args, **kwargs, interpret=self.interpret)
 
 
-class PallasInterpretModeOutOfBoundsTest(PallasTPUTest):
-
-  interpret: bool = True
-
-  def test_interpret_mode_out_of_bounds_access(self):
-    block_size = 32
-    # Create input tensors which require a reduction along an axis
-    # not divisible by block_size.
-    x = jax.random.normal(jax.random.key(0), (block_size, block_size + 1))
-    y = jax.random.normal(jax.random.key(1), (block_size + 1, block_size))
-    expected = jnp.dot(x, y)
-
-    in_specs = [
-        pl.BlockSpec(lambda i, j, k: (i, k), (block_size, block_size)),
-        pl.BlockSpec(lambda i, j, k: (k, j), (block_size, block_size)),
-    ]
-    out_spec = pl.BlockSpec(lambda i, j, k: (i, j), (block_size, block_size))
-
-    def _unmasked_matmul_kernel(x_ref, y_ref, o_ref):
-      @pl.when(pl.program_id(2) == 0)
-      def _():
-        o_ref[...] = jnp.zeros_like(o_ref)
-
-      o_ref[...] += x_ref[...] @ y_ref[...]
-
-    out = pl.pallas_call(
-        _unmasked_matmul_kernel,
-        out_shape=expected,
-        grid=(1, 1, 2),
-        in_specs=in_specs,
-        out_specs=out_spec,
-        interpret=True,
-    )(x, y)
-
-    # With a naive matmul implementation, using uninitialized values (NaN) will
-    # cause the overall output to be NaN.
-    with self.subTest('UnmaskedIsNaN'):
-      np.testing.assert_allclose(
-          np.isnan(out), jnp.ones_like(out, dtype=jnp.bool_)
-      )
-
-    def _masked_matmul_kernel(x_ref, y_ref, o_ref):
-      @pl.when(pl.program_id(2) == 0)
-      def _():
-        o_ref[:, :] = jnp.zeros_like(o_ref)
-
-      # Create a validity mask for OOB values.
-      num_valid = x.shape[1] - pl.program_id(2) * block_size
-      num_valid = jnp.minimum(num_valid, block_size)
-      mask = jnp.tril(jnp.ones_like(x_ref[:, :]))[num_valid - 1][jnp.newaxis, :]
-      mask = jnp.repeat(mask, block_size, axis=0)
-
-      # Mask and multiply.
-      masked_x = jnp.where(mask, x_ref[:, :], 0.0)
-      masked_y = jnp.where(mask.T, y_ref[:, :], 0.0)
-      o_ref[:, :] += masked_x @ masked_y
-
-    out = pl.pallas_call(
-        _masked_matmul_kernel,
-        out_shape=expected,
-        grid=(1, 1, 2),
-        in_specs=in_specs,
-        out_specs=out_spec,
-        interpret=True,
-    )(x, y)
-
-    # With a masked matmul implementation, uninitialized values will be
-    # masked before computation. This should return the correct result.
-    with self.subTest('MaskedOutputIsCorrect'):
-      np.testing.assert_allclose(out, expected, atol=1e-6)
-
-
 class PallasCallScalarPrefetchTest(PallasTPUTest):
 
   def test_trivial_scalar_prefetch(self):
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
index f35649948..abfe23b11 100644
--- a/tests/pallas/pallas_test.py
+++ b/tests/pallas/pallas_test.py
@@ -1670,29 +1670,6 @@ class PallasOpsTest(PallasTest):
 
     np.testing.assert_allclose(f(), kernel())
 
-  @parameterized.parameters("float16", "bfloat16", "float32")
-  def test_approx_tanh(self, dtype):
-    if self.INTERPRET:
-      self.skipTest("approx_tanh is not supported in interpreter mode")
-    if dtype == "bfloat16" and not self.check_gpu_capability_at_least(90):
-      self.skipTest("tanh.approx.bf16 requires a GPU with capability >= sm90")
-
-    @functools.partial(
-        self.pallas_call, out_shape=jax.ShapeDtypeStruct((4,), dtype), grid=1
-    )
-    def kernel(x_ref, o_ref):
-      o_ref[...] = plgpu.approx_tanh(x_ref[...])
-
-    x = jnp.asarray([-1, 0.42, 0.24, 1]).astype(dtype)
-    # We upcast to float32 because NumPy <2.0 does not handle custom dtypes
-    # properly. See https://github.com/google/jax/issues/11014.
-    np.testing.assert_allclose(
-        kernel(x).astype(jnp.float32),
-        jnp.tanh(x).astype(jnp.float32),
-        atol=5e-3,
-        rtol=5e-3,
-    )
-
 
 class PallasOpsInterpretTest(PallasOpsTest):
   INTERPRET = True
diff --git a/tests/sparse_nm_test.py b/tests/sparse_nm_test.py
index 602000901..12456ea16 100644
--- a/tests/sparse_nm_test.py
+++ b/tests/sparse_nm_test.py
@@ -95,8 +95,9 @@ class SpmmTest(jtu.JaxTestCase):
     jtu.check_eq(dot_sparse, dot_dense.astype(dtypes.bfloat16))
 
   # ----- Test different input types
+  # TODO(b/336519663): add int8 type once codegen is fixed
   @parameterized.product(
-      lhs_type=[jnp.int8, jnp.int16, jnp.float16, jnp.bfloat16],
+      lhs_type=[jnp.int16, jnp.float16, jnp.bfloat16],
       rhs_type=[jnp.bfloat16],
       output_type=[jnp.bfloat16, jnp.float32],
   )