Repository: jax-ml/jax
Pull request #26964 (2025-03-06 10:59:59)

AutoPGLE: force-disable graphs less
Previously, XLA's command buffers (CUDA graphs) would be disabled both for PGLE profile collection and when re-compiling using the profile data. With this change, they are only disabled when collecting the profile data.

olupton: @hawkinsp can you please take a look, or assign someone?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 68450dca4..5785f6193 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -36,14 +36,6 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
       `jax.tree_leaves`, `jax.tree_structure`, `jax.tree_transpose`, and
       `jax.tree_unflatten`. Replacements can be found in {mod}`jax.tree` or
       {mod}`jax.tree_util`.
-    * From `jax.core`: `AxisSize`, `ClosedJaxpr`, `EvalTrace`, `InDBIdx`, `InputType`,
-      `Jaxpr`, `JaxprEqn`, `Literal`, `MapPrimitive`, `OpaqueTraceState`, `OutDBIdx`,
-      `Primitive`, `Token`, `TRACER_LEAK_DEBUGGER_WARNING`, `Var`, `concrete_aval`,
-      `dedup_referents`, `escaped_tracer_error`, `extend_axis_env_nd`, `get_referent`,
-      `join_effects`, `leaked_tracer_error`, `maybe_find_leaked_tracers`, `raise_to_shaped`,
-      `raise_to_shaped_mappings`, `reset_trace_state`, `str_eqn_compact`,
-      `substitute_vars_in_output_ty`, `typecompat`, and `used_axis_names_jaxpr`. Most
-      have no public replacement, though a few are available at {mod}`jax.extend.core`.
 
 ## jax 0.5.3 (Mar 19, 2025)
 
diff --git a/benchmarks/tracing_benchmark.py b/benchmarks/tracing_benchmark.py
deleted file mode 100644
index e06ad538d..000000000
--- a/benchmarks/tracing_benchmark.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# Copyright 2025 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.
-
-"""Benchmarks for Jax tracing."""
-
-import google_benchmark
-import jax
-from jax import random
-from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel as splash
-from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
-import numpy as np
-
-
-def make_mqa_splash_attention_fn_and_args():
-  seed = 0
-  key = random.key(seed)
-  k1, k2, k3 = random.split(key, 3)
-
-  q_seq_len = 1024
-  kv_seq_len = 1024
-  num_q_heads = 2
-  head_dim_qk = 128
-  head_dim_v = 128
-  dtype = np.dtype("float32")
-
-  q = random.uniform(k1, (num_q_heads, q_seq_len, head_dim_qk), dtype=dtype)
-  k = random.uniform(k2, (kv_seq_len, head_dim_qk), dtype=dtype)
-  v = random.uniform(k3, (kv_seq_len, head_dim_v), dtype=dtype)
-
-  mask = mask_lib.NumpyMask(
-      mask_lib.make_random_mask((q_seq_len, kv_seq_len), sparsity=0.5, seed=0)
-  )
-  mask = mask_lib.MultiHeadMask(tuple(mask for _ in range(num_q_heads)))
-  block_sizes = splash.BlockSizes.get_default()
-
-  return (
-      jax.jit(
-          splash.make_splash_mqa_single_device(mask, block_sizes=block_sizes)
-      )
-  ), (q, k, v)
-
-
-@google_benchmark.register
-@google_benchmark.option.unit(google_benchmark.kMillisecond)
-def test_pallas_mqa_splash_attention_trace(state):
-  attn, (q, k, v) = make_mqa_splash_attention_fn_and_args()
-
-  while state:
-    _ = attn.trace(q, k, v)
-    jax.clear_caches()
-
-
-@google_benchmark.register
-@google_benchmark.option.unit(google_benchmark.kMillisecond)
-def test_pallas_mqa_splash_attention_lower(state):
-  attn, (q, k, v) = make_mqa_splash_attention_fn_and_args()
-  traced = attn.trace(q, k, v)
-
-  while state:
-    _ = traced.lower(lowering_platforms=("tpu",))
-    jax.clear_caches()
-
-
-if __name__ == "__main__":
-  google_benchmark.main()
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index dea532d13..9ac47aa4f 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -17,6 +17,8 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+import copy
+from functools import partial
 import logging
 import time
 from typing import Any, Callable
@@ -197,15 +199,6 @@ def get_compile_options(
       config.memory_fitting_level.value
   ).value
 
-  # This is a temporary workaround to simplify the AutoPGLE usage.
-  # TODO(b/376647494): Remove once the bug is fixed.
-  if ((config.enable_pgle.value and config.pgle_profiling_runs.value > 0)
-      or config.compilation_cache_expect_pgle.value):
-    logger.debug("Explicitly disabling command buffer scheduling for AutoPGLE.")
-    if env_options_overrides is None:
-      env_options_overrides = {}
-    env_options_overrides['xla_gpu_enable_command_buffer'] = ''
-
   if env_options_overrides is not None:
     # Some overrides are passed directly on build_options.
     overrides_on_build_options = [
@@ -298,6 +291,8 @@ def backend_compile(
     options: xc.CompileOptions,
     host_callbacks: Sequence[Any],
 ) -> xc.LoadedExecutable:
+  sym_name = module.operation.attributes['sym_name']
+  module_name = ir.StringAttr(sym_name).value
   # Convert ir.Module to a string representation, unless the backend
   # explicitly flags the ability to handle a module directly (avoiding the
   # overhead of back and forth conversions).
@@ -308,6 +303,14 @@ def backend_compile(
   else:
     built_c = module
 
+  if (options.executable_build_options.fdo_profile is not None
+      and len(options.executable_build_options.fdo_profile)):
+    logger.debug(
+        "Compiling module %s with FDO profile of length %d",
+        module_name,
+        len(options.executable_build_options.fdo_profile),
+    )
+
   try:
     # we use a separate function call to ensure that XLA compilation appears
     # separately in Python profiling results
@@ -362,72 +365,31 @@ def compile_or_get_cached(
   if dumped_to := mlir.dump_module_to_file(computation, "compile"):
     logging.info("Dumped the module to %s.", dumped_to)
 
-  use_compilation_cache = compilation_cache.is_cache_used(backend)
-
   is_multi_process = (
       len({device.process_index for device in devices.flatten()}) > 1
   )
   min_device_process_id = min(
       devices.flatten(), key=lambda device: device.id
   ).process_index
-  is_auto_pgle_used = (
-      config.enable_pgle.value and config.pgle_profiling_runs.value > 0
-  )
 
-  if not use_compilation_cache:
-    if (
-        is_multi_process
-        and is_auto_pgle_used
-        and distributed.global_state.client is not None
-    ):
-      compile_options.executable_build_options.fdo_profile = (
-          _share_fdo_profiles(
-              computation,
-              devices,
-              compile_options,
-              backend,
-              distributed.global_state.client,
-              min_device_process_id,
-          )
-      )
+  # cache_key: may be None if compilation caching is disabled
+  cache_key, compile_options = _resolve_compilation_strategy(
+    computation,
+    devices,
+    compile_options,
+    backend,
+    pgle_profiler,
+    is_multi_process,
+    module_name,
+    min_device_process_id,
+  )
 
+  if cache_key is None:
     return backend_compile(backend, computation, compile_options,
                            host_callbacks)
 
   monitoring.record_event('/jax/compilation_cache/compile_requests_use_cache')
 
-  try:
-    if config.remove_custom_partitioning_ptr_from_cache_key.value:
-      ignore_callbacks = cache_key_type.IgnoreCallbacks.CUSTOM_PARTITIONING
-    else:
-      ignore_callbacks = cache_key_type.IgnoreCallbacks.NO
-
-    cache_key = compilation_cache.get_cache_key(
-        computation,
-        devices,
-        compile_options,
-        backend,
-        ignore_callbacks=ignore_callbacks,
-    )
-  except xc._xla.XlaRuntimeError as ex:
-    logger.error("compile_or_get_cached: unable to generate cache key, "
-                 "skipping the cache: %s", ex)
-    return backend_compile(backend, computation, compile_options,
-                           host_callbacks)
-
-  if is_auto_pgle_used or config.compilation_cache_expect_pgle.value:
-    cache_key = _resolve_pgle_module_cache_key(
-        computation,
-        devices,
-        compile_options,
-        backend,
-        pgle_profiler,
-        is_multi_process,
-        cache_key,
-        module_name,
-        min_device_process_id,
-    )
-
   cache_retrieval_start = time.monotonic()
   retrieved_executable, retrieved_compile_time = _cache_read(
       module_name, cache_key, compile_options, backend)
@@ -481,85 +443,130 @@ def compile_or_get_cached(
 # 1. PGLE optimized module (the one which was recompiled with FDO profile) is
 # in the persistent cache. In this case the module should be returned from
 # cache and PGLE should be disabled for this module. Is module is stored in
-# the persistent cache under the "pgle_profiled_module_key" which calculated
-# with replacing FDO profile with flag which identify that module were PGLE
-# profiled.
+# the persistent cache under the "pgle_optimized_cache_key", which is
+# calculated by replacing the FDO profile with a sentinel value that identifies
+# that the module was optimized with PGLE.
 # 2. PGLE profiled module is not in the persistent cache and the module is
-# getting built with an FDO profile. In this case we need to share FDO profile
-# with other processes and store the result under the
-# "pgle_profiled_module_key" so later in case 1 we will be able to find the
+# getting built with an FDO profile. In this case we need to share the FDO
+# profile with any other processes and store the result under the
+# "pgle_optimized_cache_key" so later in case 1 we will be able to find the
 # module.
 # 3. PGLE profiled module is not in the persistent cache and the module is
 # getting compiled to be PGLEd (FDO profile is empty). In this case we need to
-# simply return the non-PGLE profiled module from the persistent cache.
+# simply return the non-PGLE profiled module from the persistent cache if it
+# exists, and otherwise compile it.
 #
 # If the compilation_cache_expect_pgle option is set then in case 1 the PGLE
 # optimized module will be loaded even if PGLE is not enabled in the current
 # process. This is useful if we want to combine the use of PGLE with other
 # profiling tools (e.g. Nsight Systems) that cannot co-exist with PGLE due to
 # contention for CUPTI resources.
-def _resolve_pgle_module_cache_key(
+def _resolve_compilation_strategy(
     computation: ir.Module,
     devices: np.ndarray,
     compile_options: xc.CompileOptions,
     backend: xc.Client,
     pgle_profiler: profiler.PGLEProfiler | None,
     is_multi_process: bool,
-    cache_key: str,
     module_name: str,
     min_device_process_id: int,
-) -> str:
-  fdo_profile = compile_options.executable_build_options.fdo_profile
-  compile_options.executable_build_options.fdo_profile = b"pgle profiled"
-
-  pgle_profiled_module_key = compilation_cache.get_cache_key(
-      computation,
-      devices,
-      compile_options,
-      backend,
-      cache_key_type.IgnoreCallbacks.ALL,
+) -> tuple[str | None, xc.CompileOptions]:
+  is_auto_pgle_used = (
+      config.enable_pgle.value and config.pgle_profiling_runs.value > 0
   )
-  compile_options.executable_build_options.fdo_profile = fdo_profile
-
-  result_key = cache_key
-  if _is_executable_in_cache(backend, pgle_profiled_module_key):
-    # Load PGLE profiled module from the persistent cache.
-    result_key = pgle_profiled_module_key
-    if config.compilation_cache_expect_pgle.value:
-      logging.info(f"PGLE-optimized {module_name} loaded from compilation cache")
-    if pgle_profiler is not None:
-      pgle_profiler.disable()
+
+  get_cache_key = partial(_get_cache_key, backend=backend,
+                          computation=computation, devices=devices)
+
+  if is_auto_pgle_used or config.compilation_cache_expect_pgle.value:
+    # This can be None if cache key generation fails.
+    pgle_optimized_cache_key = get_cache_key(compile_options,
+                                             override_fdo_profile=b"pgle profiled")
+    # TODO(b/376647494): remove the workaround when the bug is fixed; the JAX
+    # profiler cannot collect sufficiently detailed profile data for PGLE if
+    # command buffers / CUDA graphs are enabled. Therefore disable command
+    # buffers when compiling for PGLE data collection, but not if AutoPGLE is
+    # not enabled, and not when re-compiling using PGLE data. This condition
+    # includes `compilation_cache_expect_pgle` so that slow-to-compile modules
+    # that are not executed often enough to trigger re-compilation will still
+    # be cached between an "enable_pgle" run and an "expect_pgle" run.
+    first_pass_compile_options = copy.deepcopy(compile_options)
+    first_pass_compile_options.env_option_overrides += [
+      ("xla_gpu_enable_command_buffer", ""),
+    ]
   else:
-    # No PGLE-optimised module found in the persistent cache.
-    if (config.compilation_cache_expect_pgle.value
-        and _is_executable_in_cache(backend, cache_key)):
-      # The user asserted this miss was unexpected; emit a warning
+    pgle_optimized_cache_key = None
+    first_pass_compile_options = compile_options
+
+  # This can be None if cache key generation fails or caching is disabled
+  cache_key = get_cache_key(first_pass_compile_options)
+
+  if cache_key is not None and pgle_optimized_cache_key is not None:
+    # The compilation cache is enabled and AutoPGLE is enabled/expected
+    if _is_executable_in_cache(backend, pgle_optimized_cache_key):
+      if config.compilation_cache_expect_pgle.value:
+        logging.info(f"PGLE-optimized {module_name} loaded from compilation cache")
+      # No need to record N profiles in this case
+      if pgle_profiler is not None:
+        pgle_profiler.disable()
+      return pgle_optimized_cache_key, compile_options
+    elif (config.compilation_cache_expect_pgle.value
+          and _is_executable_in_cache(backend, cache_key)):
+      # No PGLE-optimized module found in the persistent cache, and the user
+      # asserted (expect_pgle) that this miss was unexpected
       warnings.warn(f"PERSISTENT CACHE MISS for PGLE-optimized {module_name} "
                     "despite non-PGLE hit; it may not have been executed "
                     "enough times when the cache was populated")
-    if fdo_profile is not None and len(fdo_profile) > 0:
-      # Store module under PGLE profiled module cache key.
-      result_key = pgle_profiled_module_key
-      if is_multi_process and distributed.global_state.client is not None:
-        compile_options.executable_build_options.fdo_profile = (
-            _share_fdo_profiles(
-                computation,
-                devices,
-                compile_options,
-                backend,
-                distributed.global_state.client,
-                min_device_process_id,
-            )
-        )
-      else:
-        compile_options.executable_build_options.fdo_profile = fdo_profile
-        logger.debug(
-            "Compiling module %s with FDO profile of length %d",
-            module_name,
-            len(compile_options.executable_build_options.fdo_profile),
+
+  if (is_auto_pgle_used
+      and compile_options.executable_build_options.fdo_profile is not None
+      and len(compile_options.executable_build_options.fdo_profile)):
+    # Profile data are available to trigger a PGLE-optimized recompilation;
+    # store under `pgle_optimized_cache_key` if the cache is enabled
+    if is_multi_process and distributed.global_state.client is not None:
+      compile_options.executable_build_options.fdo_profile = (
+        _share_fdo_profiles(
+            computation,
+            devices,
+            compile_options,
+            backend,
+            distributed.global_state.client,
+            min_device_process_id,
         )
-  return result_key
+      )
+    return pgle_optimized_cache_key, compile_options
+  else:
+    # Compile for PGLE collection, store under `cache_key` if the cache is
+    # enabled. This is also the AutoPGLE-disabled path.
+    return cache_key, first_pass_compile_options
 
+def _get_cache_key(
+    options: xc.CompileOptions,
+    backend: xc.Client,
+    computation: ir.Module,
+    devices: np.ndarray,
+    override_fdo_profile: bytes | None = None) -> str | None:
+  if not compilation_cache.is_cache_used(backend):
+    return None
+  if config.remove_custom_partitioning_ptr_from_cache_key.value:
+    ignore_callbacks = cache_key_type.IgnoreCallbacks.CUSTOM_PARTITIONING
+  else:
+    ignore_callbacks = cache_key_type.IgnoreCallbacks.NO
+  if override_fdo_profile is not None:
+    options = copy.deepcopy(options)
+    options.executable_build_options.fdo_profile = override_fdo_profile
+  try:
+    return compilation_cache.get_cache_key(
+        computation,
+        devices,
+        options,
+        backend,
+        ignore_callbacks,
+    )
+  except xc._xla.XlaRuntimeError as ex:
+    logger.error("compile_or_get_cached: unable to generate cache key, "
+                  "skipping the cache: %s", ex)
+  return None
 
 # The process that has the lowest device ID should share FDO profile before
 # compilation with other processes.
diff --git a/jax/_src/numpy/array_api_metadata.py b/jax/_src/numpy/array_api_metadata.py
index 5267e5121..d8d2c2d1a 100644
--- a/jax/_src/numpy/array_api_metadata.py
+++ b/jax/_src/numpy/array_api_metadata.py
@@ -24,9 +24,7 @@ from types import ModuleType
 import jax
 from jax._src.sharding import Sharding
 from jax._src.lib import xla_client as xc
-from jax._src import config
-from jax._src import dtypes as _dtypes
-from jax._src import xla_bridge as xb
+from jax._src import dtypes as _dtypes, config
 
 
 __array_api_version__ = '2023.12'
@@ -75,10 +73,7 @@ class ArrayNamespaceInfo:
     return None
 
   def devices(self):
-    out = [None]  # None indicates "uncommitted"
-    for backend in xb.backends():
-        out.extend(jax.devices(backend))
-    return out
+    return jax.devices()
 
   def capabilities(self):
     return self._capabilities
diff --git a/jax/_src/pallas/mosaic_gpu/core.py b/jax/_src/pallas/mosaic_gpu/core.py
index b0d4f23c7..8522bdf65 100644
--- a/jax/_src/pallas/mosaic_gpu/core.py
+++ b/jax/_src/pallas/mosaic_gpu/core.py
@@ -342,7 +342,7 @@ class TransposeTransform(MemoryRefTransform):
 @tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class TransposeRef(state_types.Transform):
-  permutation: tuple[int, ...] = dataclasses.field(metadata=dict(static=True))
+  permutation: tuple[int, ...]
 
   def transform_shape(self, shape):
     if shape is None:
@@ -370,30 +370,18 @@ class TransposeRef(state_types.Transform):
     return mgpu.TransposeTransform(_perm_inverse(self.permutation))
 
 
-def transform_ref(
-    ref: pallas_core.TransformedRef,
-    transform: state_types.Transform
+def transpose_ref(
+    ref: pallas_core.TransformedRef | Any,
+    permutation: tuple[int, ...],
 ) -> pallas_core.TransformedRef:
   if not isinstance(ref, pallas_core.TransformedRef):
     if not isinstance(jax_core.get_aval(ref), pallas_core.AbstractMemoryRef):
       raise TypeError("ref must be a reference")
     ref = pallas_core.TransformedRef(ref, transforms=())
   return pallas_core.TransformedRef(
-      ref.ref, (*ref.transforms, transform),
+      ref.ref, (*ref.transforms, TransposeRef(permutation)),
   )
 
-def transpose_ref(
-    ref: pallas_core.TransformedRef | Any,
-    permutation: tuple[int, ...],
-) -> pallas_core.TransformedRef:
-  return transform_ref(ref, TransposeRef(permutation))
-
-def untile_ref(ref, tiling: tuple[int, ...]) -> pallas_core.TransformedRef:
-  return transform_ref(ref, UntileRef(tiling))
-
-def unswizzle_ref(ref, swizzle: int) -> pallas_core.TransformedRef:
-  return transform_ref(ref, UnswizzleRef(swizzle))
-
 
 @dataclasses.dataclass(frozen=True)
 class SwizzleTransform(MemoryRefTransform):
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index f027d5bcb..daa718ff1 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -2369,20 +2369,20 @@ def _bitcast_convert_type_lowering_rule(
 
 @register_lowering_rule(lax.optimization_barrier_p, mgpu.ThreadSemantics.Lane)
 def _optimization_barrier_lowering(ctx: LoweringRuleContext, *args):
-  result = mgpu.optimization_barrier(
-      *(_ensure_fa(arg, aval.dtype) for arg, aval in zip(args, ctx.avals_in))
-  )
-  return (result,) if len(ctx.avals_in) == 1 else result
+  args = (_ensure_fa(arg, aval.dtype) for arg, aval in zip(args, ctx.avals_in))
+  return mgpu.optimization_barrier(*args)
 
 
 @register_lowering_rule(
     lax.optimization_barrier_p, mgpu.ThreadSemantics.Warpgroup
 )
 def _optimization_barrier_lowering_wg(ctx: LoweringRuleContext, *args):
-  result = mgpu.dialect.optimization_barrier([
+  args = [
       _ensure_ir_value(arg, aval.dtype) for arg, aval in zip(args, ctx.avals_in)
-  ])
-  return (result,) if len(ctx.avals_in) == 1 else result
+  ]
+  result = mgpu.dialect.optimization_barrier(args)
+
+  return (result,) if len(args) == 1 else result
 
 
 def _bcast(
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index 07235c2fc..d0632728f 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -759,8 +759,9 @@ def _wgmma_lowering(
       rhs_transpose = False
     case (
         gpu_core.UnswizzleRef(rhs_swizzle),
+        gpu_core.TransposeRef((1, 0, 2, 3)),  # Only transpose between tiles
         gpu_core.UntileRef(rhs_tiling),
-        gpu_core.TransposeRef((1, 0)),
+        gpu_core.TransposeRef((1, 0)),  # Transpose the two logical dims
     ):
       rhs_transpose = True
     case (
@@ -793,7 +794,7 @@ def _wgmma_lowering(
       raise NotImplementedError("WGMMA rhs tiling does not fit swizzle")
 
   if rhs_transpose:
-    b = mgpu.memref_transpose(b, (1, 0, 3, 2))
+    b = mgpu.memref_transpose(b, (0, 1, 3, 2))
   new_acc = mgpu.wgmma(acc, a, b, swizzle=rhs_swizzle)
   nvvm_dialect.wgmma_commit_group_sync_aligned()
   return new_acc
diff --git a/jax/core.py b/jax/core.py
index 688fa14d9..b404e66c2 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -81,21 +81,75 @@ from jax._src.core import (
 
 from jax._src import core as _src_core
 _deprecations = {
+    # Added 2024-12-16
+    "ClosedJaxpr": ("jax.core.ClosedJaxpr is deprecated. Use jax.extend.core.ClosedJaxpr instead, "
+                    "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+                    _src_core.ClosedJaxpr),
+    "Jaxpr": ("jax.core.Jaxpr is deprecated. Use jax.extend.core.Jaxpr instead, "
+              "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+              _src_core.Jaxpr),
+    "JaxprEqn": ("jax.core.JaxprEqn is deprecated. Use jax.extend.core.JaxprEqn instead, "
+                 "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+                 _src_core.JaxprEqn),
+    "Literal": ("jax.core.Literal is deprecated. Use jax.extend.core.Literal instead, "
+                "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+                _src_core.Literal),
+    "Primitive": ("jax.core.Primitive is deprecated. Use jax.extend.core.Primitive instead, "
+                  "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+                  _src_core.Primitive),
+    "Token": ("jax.core.Token is deprecated. Use jax.extend.core.Token instead, "
+              "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+              _src_core.Token),
+    "Var": ("jax.core.Var is deprecated. Use jax.extend.core.Var instead, "
+            "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.",
+            _src_core.Var),
     # Added 2024-12-11
     "axis_frame": ("jax.core.axis_frame is deprecated.", _src_core.axis_frame),
     "AxisName": ("jax.core.AxisName is deprecated.", _src_core.AxisName),
+    "AxisSize": ("jax.core.AxisSize is deprecated.", _src_core.AxisSize),
     "ConcretizationTypeError": ("jax.core.ConcretizationTypeError is deprecated; "
                                 "use jax.errors.ConcretizationTypeError.",
                                 _src_core.ConcretizationTypeError),
+    "EvalTrace": ("jax.core.EvalTrace is deprecated.", _src_core.EvalTrace),
+    "InDBIdx": ("jax.core.InDBIdx is deprecated.", _src_core.InDBIdx),
+    "InputType": ("jax.core.InputType is deprecated.", _src_core.InputType),
+    "MapPrimitive": ("jax.core.MapPrimitive is deprecated.", _src_core.MapPrimitive),
+    "OpaqueTraceState": ("jax.core.OpaqueTraceState is deprecated.", _src_core.OpaqueTraceState),
+    "OutDBIdx": ("jax.core.OutDBIdx is deprecated.", _src_core.OutDBIdx),
+    "TRACER_LEAK_DEBUGGER_WARNING": ("jax.core.TRACER_LEAK_DEBUGGER_WARNING is deprecated.",
+                                     _src_core.TRACER_LEAK_DEBUGGER_WARNING),
     "call_p": ("jax.core.call_p is deprecated. Use jax.extend.core.primitives.call_p",
                _src_core.call_p),
     "closed_call_p": ("jax.core.closed_call_p is deprecated. Use jax.extend.core.primitives.closed_call_p",
                       _src_core.closed_call_p),
+    "concrete_aval": ("jax.core.concrete_aval is deprecated.", _src_core.abstractify),
+    "dedup_referents": ("jax.core.dedup_referents is deprecated.", _src_core.dedup_referents),
+    "escaped_tracer_error": ("jax.core.escaped_tracer_error is deprecated.",
+                             _src_core.escaped_tracer_error),
+    "extend_axis_env_nd": ("jax.core.extend_axis_env_nd is deprecated.",
+                           _src_core.extend_axis_env_nd),
     "get_type": ("jax.core.get_type is deprecated.", _src_core.get_aval),
+    "get_referent": ("jax.core.get_referent is deprecated.", _src_core.get_referent),
+    "join_effects": ("jax.core.join_effects is deprecated.", _src_core.join_effects),
+    "leaked_tracer_error": ("jax.core.leaked_tracer_error is deprecated.",
+                            _src_core.leaked_tracer_error),
+    "maybe_find_leaked_tracers": ("jax.core.maybe_find_leaked_tracers is deprecated.",
+                                  _src_core.maybe_find_leaked_tracers),
+    "raise_to_shaped_mappings": ("jax.core.raise_to_shaped_mappings is deprecated."
+                                 " It is unused as of jax v0.4.36.",
+                                 _src_core.raise_to_shaped_mappings),
+    "reset_trace_state": ("jax.core.reset_trace_state is deprecated.",
+                          _src_core.reset_trace_state),
+    "str_eqn_compact": ("jax.core.str_eqn_compact is deprecated.", _src_core.str_eqn_compact),
+    "substitute_vars_in_output_ty": ("jax.core.substitute_vars_in_output_ty is deprecated.",
+                                     _src_core.substitute_vars_in_output_ty),
     "trace_state_clean": ("jax.core.trace_state_clean is deprecated.",
                           _src_core.trace_state_clean),
     "typecheck": ("jax.core.typecheck is deprecated.", _src_core.typecheck),
+    "typecompat": ("jax.core.typecompat is deprecated.", _src_core.typecompat),
     "typematch": ("jax.core.typematch is deprecated.", _src_core.typematch),
+    "used_axis_names_jaxpr": ("jax.core.used_axis_names_jaxpr is deprecated.",
+                              _src_core.used_axis_names_jaxpr),
     # Added 2024-12-10
     "full_lower": ("jax.core.full_lower is deprecated. It is a no-op as of JAX v0.4.36.",
                    _src_core.full_lower),
@@ -104,61 +158,54 @@ _deprecations = {
                      _src_core.jaxpr_as_fun),
     "lattice_join": ("jax.core.lattice_join is deprecated. It is a no-op as of JAX v0.4.36.",
                      _src_core.lattice_join),
-    # Finalized 2025-03-25 for JAX v0.6.0; remove after 2025-06-25
-    "AxisSize": ("jax.core.AxisSize was removed in JAX v0.6.0.", None),
-    "ClosedJaxpr": ("jax.core.ClosedJaxpr was removed in JAX v0.6.0. Use jax.extend.core.ClosedJaxpr instead, "
-                    "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "EvalTrace": ("jax.core.EvalTrace was removed in JAX v0.6.0.", None),
-    "InDBIdx": ("jax.core.InDBIdx was removed in JAX v0.6.0.", None),
-    "InputType": ("jax.core.InputType was removed in JAX v0.6.0.", None),
-    "Jaxpr": ("jax.core.Jaxpr was removed in JAX v0.6.0. Use jax.extend.core.Jaxpr instead, "
-              "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "JaxprEqn": ("jax.core.JaxprEqn was removed in JAX v0.6.0. Use jax.extend.core.JaxprEqn instead, "
-                 "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "Literal": ("jax.core.Literal was removed in JAX v0.6.0. Use jax.extend.core.Literal instead, "
-                "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "MapPrimitive": ("jax.core.MapPrimitive was removed in JAX v0.6.0.", None),
-    "OpaqueTraceState": ("jax.core.OpaqueTraceState was removed in JAX v0.6.0.", None),
-    "OutDBIdx": ("jax.core.OutDBIdx was removed in JAX v0.6.0.", None),
-    "Primitive": ("jax.core.Primitive was removed in JAX v0.6.0. Use jax.extend.core.Primitive instead, "
-                  "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "Token": ("jax.core.Token was removed in JAX v0.6.0. Use jax.extend.core.Token instead, "
-              "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "TRACER_LEAK_DEBUGGER_WARNING": ("jax.core.TRACER_LEAK_DEBUGGER_WARNING was removed in JAX v0.6.0.", None),
-    "Var": ("jax.core.Var was removed in JAX v0.6.0. Use jax.extend.core.Var instead, "
-            "and see https://jax.readthedocs.io/en/latest/jax.extend.html for details.", None),
-    "concrete_aval": ("jax.core.concrete_aval was removed in JAX v0.6.0.", None),
-    "dedup_referents": ("jax.core.dedup_referents was removed in JAX v0.6.0.", None),
-    "escaped_tracer_error": ("jax.core.escaped_tracer_error was removed in JAX v0.6.0.", None),
-    "extend_axis_env_nd": ("jax.core.extend_axis_env_nd was removed in JAX v0.6.0.", None),
-    "get_referent": ("jax.core.get_referent was removed in JAX v0.6.0.", None),
-    "join_effects": ("jax.core.join_effects was removed in JAX v0.6.0.", None),
-    "leaked_tracer_error": ("jax.core.leaked_tracer_error was removed in JAX v0.6.0.", None),
-    "maybe_find_leaked_tracers": ("jax.core.maybe_find_leaked_tracers was removed in JAX v0.6.0.", None),
-    "raise_to_shaped": ("jax.core.raise_to_shaped was removed in JAX v0.6.0. It is a no-op as of JAX v0.4.36.", None),
-    "raise_to_shaped_mappings": ("jax.core.raise_to_shaped_mappings was removed in JAX v0.6.0."
-                                 " It is unused as of jax v0.4.36.", None),
-    "reset_trace_state": ("jax.core.reset_trace_state was removed in JAX v0.6.0.", None),
-    "str_eqn_compact": ("jax.core.str_eqn_compact was removed in JAX v0.6.0.", None),
-    "substitute_vars_in_output_ty": ("jax.core.substitute_vars_in_output_ty was removed in JAX v0.6.0.", None),
-    "typecompat": ("jax.core.typecompat was removed in JAX v0.6.0.", None),
-    "used_axis_names_jaxpr": ("jax.core.used_axis_names_jaxpr was removed in JAX v0.6.0.", None),
+    "raise_to_shaped": ("jax.core.raise_to_shaped is deprecated. It is a no-op as of JAX v0.4.36.",
+                        _src_core.raise_to_shaped),
 }
 
 import typing
 if typing.TYPE_CHECKING:
   AxisName = _src_core.AxisName
+  AxisSize = _src_core.AxisSize
+  ClosedJaxpr = _src_core.ClosedJaxpr
   ConcretizationTypeError = _src_core.ConcretizationTypeError
+  EvalTrace = _src_core.EvalTrace
+  InDBIdx = _src_core.InDBIdx
+  InputType = _src_core.InputType
+  Jaxpr = _src_core.Jaxpr
+  JaxprEqn = _src_core.JaxprEqn
+  Literal = _src_core.Literal
+  MapPrimitive = _src_core.MapPrimitive
+  OpaqueTraceState = _src_core.OpaqueTraceState
+  OutDBIdx = _src_core.OutDBIdx
+  Primitive = _src_core.Primitive
+  Token = _src_core.Token
+  TRACER_LEAK_DEBUGGER_WARNING = _src_core.TRACER_LEAK_DEBUGGER_WARNING
+  Var = _src_core.Var
   axis_frame = _src_core.axis_frame
   call_p = _src_core.call_p
   closed_call_p = _src_core.closed_call_p
+  concrete_aval = _src_core.abstractify
+  dedup_referents = _src_core.dedup_referents
+  escaped_tracer_error = _src_core.escaped_tracer_error
+  extend_axis_env_nd = _src_core.extend_axis_env_nd
   full_lower = _src_core.full_lower
   get_type = _src_core.get_aval
+  get_referent = _src_core.get_referent
   jaxpr_as_fun = _src_core.jaxpr_as_fun
+  join_effects = _src_core.join_effects
   lattice_join = _src_core.lattice_join
+  leaked_tracer_error = _src_core.leaked_tracer_error
+  maybe_find_leaked_tracers = _src_core.maybe_find_leaked_tracers
+  raise_to_shaped = _src_core.raise_to_shaped
+  raise_to_shaped_mappings = _src_core.raise_to_shaped_mappings
+  reset_trace_state = _src_core.reset_trace_state
+  str_eqn_compact = _src_core.str_eqn_compact
+  substitute_vars_in_output_ty = _src_core.substitute_vars_in_output_ty
   trace_state_clean = _src_core.trace_state_clean
   typecheck = _src_core.typecheck
+  typecompat = _src_core.typecompat
   typematch = _src_core.typematch
+  used_axis_names_jaxpr = _src_core.used_axis_names_jaxpr
 else:
   from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
   __getattr__ = _deprecation_getattr(__name__, _deprecations)
diff --git a/jax/experimental/pallas/mosaic_gpu.py b/jax/experimental/pallas/mosaic_gpu.py
index b44c86ea7..e4c5ffe04 100644
--- a/jax/experimental/pallas/mosaic_gpu.py
+++ b/jax/experimental/pallas/mosaic_gpu.py
@@ -27,10 +27,7 @@ from jax._src.pallas.mosaic_gpu.core import kernel as kernel
 from jax._src.pallas.mosaic_gpu.core import SemaphoreType as SemaphoreType
 from jax._src.pallas.mosaic_gpu.core import SwizzleTransform as SwizzleTransform
 from jax._src.pallas.mosaic_gpu.core import TilingTransform as TilingTransform
-from jax._src.pallas.mosaic_gpu.core import transform_ref as transform_ref
 from jax._src.pallas.mosaic_gpu.core import transpose_ref as transpose_ref
-from jax._src.pallas.mosaic_gpu.core import untile_ref as untile_ref
-from jax._src.pallas.mosaic_gpu.core import unswizzle_ref as unswizzle_ref
 from jax._src.pallas.mosaic_gpu.core import TransposeTransform as TransposeTransform
 from jax._src.pallas.mosaic_gpu.core import WGMMAAccumulatorRef as ACC  # noqa: F401
 from jax._src.pallas.mosaic_gpu.core import WGMMAAccumulatorRef as WGMMAAccumulatorRef
diff --git a/jax/experimental/pallas/ops/gpu/attention_mgpu.py b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
index b19e371a1..48a0d1845 100644
--- a/jax/experimental/pallas/ops/gpu/attention_mgpu.py
+++ b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
@@ -201,7 +201,7 @@ def attention(q, k, v, config: TuningConfig):
     )
     k_scratch = plgpu.SMEM(
         (max_concurrent_steps, block_kv, head_dim), jnp.float16,
-        transforms=(tiling, swizzle),
+        transforms=(tiling, plgpu.TransposeTransform((0, 2, 1, 3, 4)), swizzle),
     )
     v_scratch = plgpu.SMEM(
         (max_concurrent_steps, block_kv, head_dim), jnp.float16,
@@ -265,6 +265,7 @@ def attention_with_pipeline_emitter(q, k, v, config: TuningConfig):
 
   tiling = plgpu.TilingTransform((8, 64))
   swizzle = plgpu.SwizzleTransform(128)
+  transpose = plgpu.TransposeTransform((0, 2, 1, 3, 4))
 
   def fa3_kernel(q_ref, k_ref, v_ref, out_ref, scoped):
     batch = lax.axis_index("batch")
@@ -353,7 +354,7 @@ def attention_with_pipeline_emitter(q, k, v, config: TuningConfig):
             plgpu.GPUBlockSpec(  # k
                 block_shape=(block_kv, head_dim),
                 index_map=lambda i: (i, 0),
-                transforms=[tiling, swizzle]),
+                transforms=[tiling, transpose, swizzle]),
             plgpu.GPUBlockSpec(  # v
                 block_shape=(block_kv, head_dim),
                 index_map=lambda i: (i, 0),
diff --git a/jaxlib/cuda/BUILD b/jaxlib/cuda/BUILD
index d35e421ef..fac62c81d 100644
--- a/jaxlib/cuda/BUILD
+++ b/jaxlib/cuda/BUILD
@@ -689,7 +689,6 @@ cc_library(
         "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
         "@xla//xla:comparison_util",
         "@xla//xla:shape_util",
-        "@xla//xla:util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/ffi:ffi_api",
         "@xla//xla/ffi/api:ffi",
diff --git a/jaxlib/gpu/blas.cc b/jaxlib/gpu/blas.cc
index 59bf2c460..cf391e07e 100644
--- a/jaxlib/gpu/blas.cc
+++ b/jaxlib/gpu/blas.cc
@@ -49,6 +49,14 @@ BlasType DtypeToBlasType(const dtype& np_type) {
   return it->second;
 }
 
+// Returns the descriptor for a GetrfBatched operation.
+std::pair<size_t, nb::bytes> BuildGetrfBatchedDescriptor(const dtype& dtype,
+                                                         int b, int n) {
+  BlasType type = DtypeToBlasType(dtype);
+  size_t size = b * sizeof(void*);
+  return {size, PackDescriptor(GetrfBatchedDescriptor{type, b, n})};
+}
+
 // Returns the descriptor for a GetrfBatched operation.
 std::pair<size_t, nb::bytes> BuildGeqrfBatchedDescriptor(const dtype& dtype,
                                                          int b, int m, int n) {
@@ -59,6 +67,7 @@ std::pair<size_t, nb::bytes> BuildGeqrfBatchedDescriptor(const dtype& dtype,
 
 nb::dict Registrations() {
   nb::dict dict;
+  dict[JAX_GPU_PREFIX "blas_getrf_batched"] = EncapsulateFunction(GetrfBatched);
   dict[JAX_GPU_PREFIX "blas_geqrf_batched"] = EncapsulateFunction(GeqrfBatched);
   return dict;
 }
@@ -67,6 +76,7 @@ NB_MODULE(_blas, m) {
   tsl::ImportNumpy();
 
   m.def("registrations", &Registrations);
+  m.def("build_getrf_batched_descriptor", &BuildGetrfBatchedDescriptor);
   m.def("build_geqrf_batched_descriptor", &BuildGeqrfBatchedDescriptor);
 }
 
diff --git a/jaxlib/gpu/blas_kernels.cc b/jaxlib/gpu/blas_kernels.cc
index cdcc154d0..ac30aa9cc 100644
--- a/jaxlib/gpu/blas_kernels.cc
+++ b/jaxlib/gpu/blas_kernels.cc
@@ -52,6 +52,66 @@ int SizeOfBlasType(BlasType type) {
 
 }  // namespace
 
+// Batched LU decomposition: getrfbatched
+
+static absl::Status GetrfBatched_(gpuStream_t stream, void** buffers,
+                                  const char* opaque, size_t opaque_len) {
+  auto s = UnpackDescriptor<GetrfBatchedDescriptor>(opaque, opaque_len);
+  JAX_RETURN_IF_ERROR(s.status());
+  const GetrfBatchedDescriptor& d = **s;
+  auto h = BlasHandlePool::Borrow(stream);
+  JAX_RETURN_IF_ERROR(h.status());
+  auto& handle = *h;
+  if (buffers[0] != buffers[1]) {
+    JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuMemcpyAsync(
+        buffers[1], buffers[0], SizeOfBlasType(d.type) * d.batch * d.n * d.n,
+        gpuMemcpyDeviceToDevice, stream)));
+  }
+
+  int* ipiv = static_cast<int*>(buffers[2]);
+  int* info = static_cast<int*>(buffers[3]);
+  MakeBatchPointersAsync(stream, buffers[1], buffers[4], d.batch,
+                         SizeOfBlasType(d.type) * d.n * d.n);
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuGetLastError()));
+  switch (d.type) {
+    case BlasType::F32: {
+      float** batch_ptrs = static_cast<float**>(buffers[4]);
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasSgetrfBatched(
+          handle.get(), d.n, batch_ptrs, d.n, ipiv, info, d.batch)));
+      break;
+    }
+    case BlasType::F64: {
+      double** batch_ptrs = static_cast<double**>(buffers[4]);
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasDgetrfBatched(
+          handle.get(), d.n, batch_ptrs, d.n, ipiv, info, d.batch)));
+      break;
+    }
+    case BlasType::C64: {
+      gpublasComplex** batch_ptrs = static_cast<gpublasComplex**>(buffers[4]);
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasCgetrfBatched(
+          handle.get(), d.n, batch_ptrs, d.n, ipiv, info, d.batch)));
+      break;
+    }
+    case BlasType::C128: {
+      gpublasDoubleComplex** batch_ptrs =
+          static_cast<gpublasDoubleComplex**>(buffers[4]);
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasZgetrfBatched(
+          handle.get(), d.n, batch_ptrs, d.n, ipiv, info, d.batch)));
+      break;
+    }
+  }
+  return absl::OkStatus();
+}
+
+void GetrfBatched(gpuStream_t stream, void** buffers, const char* opaque,
+                  size_t opaque_len, XlaCustomCallStatus* status) {
+  auto s = GetrfBatched_(stream, buffers, opaque, opaque_len);
+  if (!s.ok()) {
+    XlaCustomCallStatusSetFailure(status, std::string(s.message()).c_str(),
+                                  s.message().length());
+  }
+}
+
 // Batched QR decomposition: geqrfbatched
 
 static absl::Status GeqrfBatched_(gpuStream_t stream, void** buffers,
diff --git a/jaxlib/gpu/blas_kernels.h b/jaxlib/gpu/blas_kernels.h
index 8ca7b4db4..724565ea7 100644
--- a/jaxlib/gpu/blas_kernels.h
+++ b/jaxlib/gpu/blas_kernels.h
@@ -32,6 +32,16 @@ enum class BlasType {
   C128,
 };
 
+// Batched LU decomposition: getrfbatched
+
+struct GetrfBatchedDescriptor {
+  BlasType type;
+  int batch, n;
+};
+
+void GetrfBatched(gpuStream_t stream, void** buffers, const char* opaque,
+                  size_t opaque_len, XlaCustomCallStatus* status);
+
 // Batched QR decomposition: geqrfbatched
 
 struct GeqrfBatchedDescriptor {
diff --git a/jaxlib/gpu/gpu_kernels.cc b/jaxlib/gpu/gpu_kernels.cc
index 840c313f2..242078357 100644
--- a/jaxlib/gpu/gpu_kernels.cc
+++ b/jaxlib/gpu/gpu_kernels.cc
@@ -33,10 +33,13 @@ namespace jax {
 namespace JAX_GPU_NAMESPACE {
 namespace {
 
+XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cublas_getrf_batched", GetrfBatched,
+                                         "CUDA");
 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cublas_geqrf_batched", GeqrfBatched,
                                          "CUDA");
 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cudnn_rnn", RNNForward, "CUDA");
 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cudnn_rnn_bwd", RNNBackward, "CUDA");
+XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM("cusolver_getrf", Getrf, "CUDA");
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_getrf_ffi", "CUDA",
                          GetrfFfi);
 XLA_FFI_REGISTER_HANDLER(XLA_FFI_GetApi(), "cusolver_syrk_ffi", "CUDA",
diff --git a/jaxlib/gpu/py_client_gpu.cc b/jaxlib/gpu/py_client_gpu.cc
index 38f2ac189..861ffce3e 100644
--- a/jaxlib/gpu/py_client_gpu.cc
+++ b/jaxlib/gpu/py_client_gpu.cc
@@ -43,7 +43,6 @@ limitations under the License.
 #include "xla/python/nb_numpy.h"
 #include "xla/python/types.h"
 #include "xla/shape_util.h"
-#include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
 namespace nb = nanobind;
@@ -81,14 +80,13 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
   for (size_t i = 0; i < arity; ++i) {
     auto arg = args.get<xla::ffi::AnyBuffer>(i);
     auto ptype = static_cast<xla::PrimitiveType>(arg->element_type());
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    if (ptype == xla::S1 || ptype == xla::U1) {
+    // TODO(b/395428868): Remove this check once we support subbyte types.
+    if (ptype == xla::S1 || ptype == xla::S2 || ptype == xla::S4 ||
+        ptype == xla::U1 || ptype == xla::U2 || ptype == xla::U4) {
       return xla::ffi::Error(xla::ffi::ErrorCode::kUnimplemented,
                              absl::StrFormat("Unsupported primitive type: %s",
                                              PrimitiveType_Name(ptype)));
     }
-
     if (ptype == xla::TOKEN) {
       host_input_buffers[i] = nullptr;
       continue;
@@ -114,6 +112,9 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
       PyTuple_SET_ITEM(host_input_arrays.ptr(), i, nb::none().inc_ref().ptr());
       continue;
     }
+    nb::capsule base(host_input_buffers[i], [](void* ptr) noexcept {
+      delete[] static_cast<char*>(ptr);
+    });
     auto maybe_dtype = PrimitiveTypeToNbDtype(ptype);
     if (!maybe_dtype.ok()) {
       return xla::ffi::Error::Internal(maybe_dtype.status().ToString());
@@ -121,22 +122,8 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
     auto dtype = maybe_dtype.value();
     auto dims = absl::Span<const int64_t>(arg->dimensions().begin(),
                                           arg->dimensions().size());
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    // We pass in data using default numpy layout i.e., std::nullopt.
-    size_t bits_per_element = xla::primitive_util::BitWidth(ptype);
-    if (bits_per_element == 2 || bits_per_element == 4) {
-      auto buffer = xla::UnpackIntN(
-          bits_per_element, static_cast<const char*>(host_input_buffers[i]),
-          arg->size_bytes());
-      delete[] static_cast<char*>(host_input_buffers[i]);
-      host_input_buffers[i] = buffer.release();
-    }
-    nb::capsule base(host_input_buffers[i], [](void* ptr) noexcept {
-      delete[] static_cast<char*>(ptr);
-    });
     auto array = xla::nb_numpy_ndarray(dtype, dims, std::nullopt,
-                                       host_input_buffers[i], /*base=*/base);
+                                       host_input_buffers[i], base);
     array.attr("flags").attr("writeable") = nb::bool_(false);
     PyTuple_SET_ITEM(host_input_arrays.ptr(), i, array.inc_ref().ptr());
   }
@@ -159,7 +146,8 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
     auto ret = rets.get<xla::ffi::AnyBuffer>(i).value();
     auto ptype = static_cast<xla::PrimitiveType>(ret->element_type());
     // TODO(b/395428868): Remove this check once we support subbyte types.
-    if (ptype == xla::S1 || ptype == xla::U1) {
+    if (ptype == xla::S1 || ptype == xla::S2 || ptype == xla::S4 ||
+        ptype == xla::U1 || ptype == xla::U2 || ptype == xla::U4) {
       return xla::ffi::Error(xla::ffi::ErrorCode::kUnimplemented,
                              absl::StrFormat("Unsupported primitive type: %s",
                                              PrimitiveType_Name(ptype)));
@@ -180,43 +168,32 @@ xla::ffi::Error XlaFfiPythonGpuCallback(gpuStream_t stream,
     }
     auto expected_shape = maybe_expected_shape.value();
     auto expected_strides = xla::ByteStridesForShape(expected_shape);
-
-    const void* data = array.data();
-    if (strides != expected_strides) {
-      xla::TransposePlan::Options options;
-      options.elem_size_in_bytes = xla::primitive_util::ByteWidth(ptype);
-      options.dims = absl::Span<int64_t const>(
-          reinterpret_cast<const int64_t*>(array.shape()), array.ndim());
-      absl::InlinedVector<int64_t, 4> reversed_layout;
-      reversed_layout.resize(expected_shape.dimensions().size());
-      absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
-                           reversed_layout.begin());
-      options.permutation = reversed_layout;
-      options.input_layout = xla::TransposePlan::Striding{strides};
-      auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
-      if (!maybe_plan.ok()) {
-        return xla::ffi::Error::Internal(maybe_plan.status().ToString());
-      }
-      auto plan = maybe_plan.value();
-      void* temp = new char[ret->size_bytes()];
-      temp_buffers.push_back(temp);
-      plan->Execute(data, temp);
-      data = temp;
+    if (strides == expected_strides) {
+      auto gpu_res =
+          gpuMemcpyAsync(ret->untyped_data(), array.data(), ret->size_bytes(),
+                         gpuMemcpyHostToDevice, stream);
+      CHECK_EQ(gpu_res, gpuSuccess) << "Failed to gpuMemcpyAsync";
+      continue;
     }
-
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    std::unique_ptr<char[]> buffer;
-    size_t bits_per_element = xla::primitive_util::BitWidth(ptype);
-    if (bits_per_element == 2 || bits_per_element == 4) {
-      // NOTE(dsuo): FFI arguments and return buffers are sized assuming
-      // minimum 1-byte element sizes, even if the data itself is packed.
-      buffer = xla::PackIntN(bits_per_element, static_cast<const char*>(data),
-                             ret->size_bytes());
-      data = buffer.get();
+    void* temp = new char[ret->size_bytes()];
+    temp_buffers.push_back(temp);
+    xla::TransposePlan::Options options;
+    options.elem_size_in_bytes = xla::primitive_util::ByteWidth(ptype);
+    options.dims = absl::Span<int64_t const>(
+        reinterpret_cast<const int64_t*>(array.shape()), array.ndim());
+    absl::InlinedVector<int64_t, 4> reversed_layout;
+    reversed_layout.resize(expected_shape.dimensions().size());
+    absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
+                         reversed_layout.begin());
+    options.permutation = reversed_layout;
+    options.input_layout = xla::TransposePlan::Striding{strides};
+    auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
+    if (!maybe_plan.ok()) {
+      return xla::ffi::Error::Internal(maybe_plan.status().ToString());
     }
-
-    auto gpu_res = gpuMemcpyAsync(ret->untyped_data(), data, ret->size_bytes(),
+    auto plan = maybe_plan.value();
+    plan->Execute(array.data(), temp);
+    auto gpu_res = gpuMemcpyAsync(ret->untyped_data(), temp, ret->size_bytes(),
                                   gpuMemcpyHostToDevice, stream);
     CHECK_EQ(gpu_res, gpuSuccess) << "Failed to gpuMemcpyAsync";
   }
diff --git a/jaxlib/gpu/solver.cc b/jaxlib/gpu/solver.cc
index 8013d9877..20fc30810 100644
--- a/jaxlib/gpu/solver.cc
+++ b/jaxlib/gpu/solver.cc
@@ -54,6 +54,45 @@ SolverType DtypeToSolverType(const dtype& np_type) {
   return it->second;
 }
 
+// getrf: LU decomposition
+
+// Returns the workspace size and a descriptor for a getrf operation.
+std::pair<int, nb::bytes> BuildGetrfDescriptor(const dtype& dtype, int b, int m,
+                                               int n) {
+  SolverType type = DtypeToSolverType(dtype);
+  auto h = SolverHandlePool::Borrow(/*stream=*/nullptr);
+  JAX_THROW_IF_ERROR(h.status());
+  auto& handle = *h;
+  int lwork;
+  switch (type) {
+    case SolverType::F32:
+      JAX_THROW_IF_ERROR(
+          JAX_AS_STATUS(gpusolverDnSgetrf_bufferSize(handle.get(), m, n,
+                                                     /*A=*/nullptr,
+                                                     /*lda=*/m, &lwork)));
+      break;
+    case SolverType::F64:
+      JAX_THROW_IF_ERROR(
+          JAX_AS_STATUS(gpusolverDnDgetrf_bufferSize(handle.get(), m, n,
+                                                     /*A=*/nullptr,
+                                                     /*lda=*/m, &lwork)));
+      break;
+    case SolverType::C64:
+      JAX_THROW_IF_ERROR(
+          JAX_AS_STATUS(gpusolverDnCgetrf_bufferSize(handle.get(), m, n,
+                                                     /*A=*/nullptr,
+                                                     /*lda=*/m, &lwork)));
+      break;
+    case SolverType::C128:
+      JAX_THROW_IF_ERROR(
+          JAX_AS_STATUS(gpusolverDnZgetrf_bufferSize(handle.get(), m, n,
+                                                     /*A=*/nullptr,
+                                                     /*lda=*/m, &lwork)));
+      break;
+  }
+  return {lwork, PackDescriptor(GetrfDescriptor{type, b, m, n, lwork})};
+}
+
 // geqrf: QR decomposition
 
 // Returns the workspace size and a descriptor for a geqrf operation.
@@ -423,6 +462,7 @@ std::pair<int, nb::bytes> BuildSytrdDescriptor(const dtype& dtype, bool lower,
 
 nb::dict Registrations() {
   nb::dict dict;
+  dict[JAX_GPU_PREFIX "solver_getrf"] = EncapsulateFunction(Getrf);
   dict[JAX_GPU_PREFIX "solver_geqrf"] = EncapsulateFunction(Geqrf);
   dict[JAX_GPU_PREFIX "solver_orgqr"] = EncapsulateFunction(Orgqr);
   dict[JAX_GPU_PREFIX "solver_syevd"] = EncapsulateFunction(Syevd);
@@ -456,6 +496,7 @@ nb::dict Registrations() {
 NB_MODULE(_solver, m) {
   tsl::ImportNumpy();
   m.def("registrations", &Registrations);
+  m.def("build_getrf_descriptor", &BuildGetrfDescriptor);
   m.def("build_geqrf_descriptor", &BuildGeqrfDescriptor);
   m.def("build_orgqr_descriptor", &BuildOrgqrDescriptor);
   m.def("build_syevd_descriptor", &BuildSyevdDescriptor);
diff --git a/jaxlib/gpu/solver_kernels.cc b/jaxlib/gpu/solver_kernels.cc
index 8971619d7..8c22dfcdb 100644
--- a/jaxlib/gpu/solver_kernels.cc
+++ b/jaxlib/gpu/solver_kernels.cc
@@ -50,6 +50,89 @@ static int SizeOfSolverType(SolverType type) {
   }
 }
 
+// getrf: LU decomposition
+
+static absl::Status Getrf_(gpuStream_t stream, void** buffers,
+                           const char* opaque, size_t opaque_len) {
+  auto s = UnpackDescriptor<GetrfDescriptor>(opaque, opaque_len);
+  JAX_RETURN_IF_ERROR(s.status());
+  const GetrfDescriptor& d = **s;
+  auto h = SolverHandlePool::Borrow(stream);
+  JAX_RETURN_IF_ERROR(h.status());
+  auto& handle = *h;
+  if (buffers[1] != buffers[0]) {
+    JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuMemcpyAsync(
+        buffers[1], buffers[0],
+        SizeOfSolverType(d.type) * static_cast<std::int64_t>(d.batch) *
+            static_cast<std::int64_t>(d.m) * static_cast<std::int64_t>(d.n),
+        gpuMemcpyDeviceToDevice, stream)));
+  }
+
+  int* ipiv = static_cast<int*>(buffers[2]);
+  int* info = static_cast<int*>(buffers[3]);
+  void* workspace = buffers[4];
+  switch (d.type) {
+    case SolverType::F32: {
+      float* a = static_cast<float*>(buffers[1]);
+      for (int i = 0; i < d.batch; ++i) {
+        JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnSgetrf(
+            handle.get(), d.m, d.n, a, d.m, static_cast<float*>(workspace),
+            d.lwork, ipiv, info)));
+        a += d.m * d.n;
+        ipiv += std::min(d.m, d.n);
+        ++info;
+      }
+      break;
+    }
+    case SolverType::F64: {
+      double* a = static_cast<double*>(buffers[1]);
+      for (int i = 0; i < d.batch; ++i) {
+        JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnDgetrf(
+            handle.get(), d.m, d.n, a, d.m, static_cast<double*>(workspace),
+            d.lwork, ipiv, info)));
+        a += d.m * d.n;
+        ipiv += std::min(d.m, d.n);
+        ++info;
+      }
+      break;
+    }
+    case SolverType::C64: {
+      gpuComplex* a = static_cast<gpuComplex*>(buffers[1]);
+      for (int i = 0; i < d.batch; ++i) {
+        JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnCgetrf(
+            handle.get(), d.m, d.n, a, d.m, static_cast<gpuComplex*>(workspace),
+            d.lwork, ipiv, info)));
+        a += d.m * d.n;
+        ipiv += std::min(d.m, d.n);
+        ++info;
+      }
+      break;
+    }
+    case SolverType::C128: {
+      gpuDoubleComplex* a = static_cast<gpuDoubleComplex*>(buffers[1]);
+      for (int i = 0; i < d.batch; ++i) {
+        JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnZgetrf(
+            handle.get(), d.m, d.n, a, d.m,
+            static_cast<gpuDoubleComplex*>(workspace), d.lwork, ipiv, info)));
+        a += d.m * d.n;
+        ipiv += std::min(d.m, d.n);
+        ++info;
+      }
+      break;
+    }
+  }
+  return absl::OkStatus();
+}
+
+void Getrf(gpuStream_t stream, void** buffers, const char* opaque,
+           size_t opaque_len, XlaCustomCallStatus* status) {
+  auto s = Getrf_(stream, buffers, opaque, opaque_len);
+  if (!s.ok()) {
+    XlaCustomCallStatusSetFailure(status, std::string(s.message()).c_str(),
+                                  s.message().length());
+  }
+}
+
 // geqrf: QR decomposition
 
 static absl::Status Geqrf_(gpuStream_t stream, void** buffers,
diff --git a/jaxlib/gpu/solver_kernels.h b/jaxlib/gpu/solver_kernels.h
index 6372e55b9..51082f2fe 100644
--- a/jaxlib/gpu/solver_kernels.h
+++ b/jaxlib/gpu/solver_kernels.h
@@ -33,6 +33,16 @@ enum class SolverType {
   C128,
 };
 
+// getrf: LU decomposition
+
+struct GetrfDescriptor {
+  SolverType type;
+  int batch, m, n, lwork;
+};
+
+void Getrf(gpuStream_t stream, void** buffers, const char* opaque,
+           size_t opaque_len, XlaCustomCallStatus* status);
+
 // geqrf: QR decomposition
 
 struct GeqrfDescriptor {
diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD
index 358a6d1cc..d0c0c798a 100644
--- a/jaxlib/rocm/BUILD
+++ b/jaxlib/rocm/BUILD
@@ -588,7 +588,6 @@ cc_library(
         "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
         "@xla//xla:comparison_util",
         "@xla//xla:shape_util",
-        "@xla//xla:util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/ffi:ffi_api",
         "@xla//xla/ffi/api:ffi",
diff --git a/jaxlib/xla/BUILD b/jaxlib/xla/BUILD
index 5b532c1dc..2ca18afda 100644
--- a/jaxlib/xla/BUILD
+++ b/jaxlib/xla/BUILD
@@ -637,7 +637,6 @@ cc_library(
         "@nanobind",
         "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
         "@xla//xla:shape_util",
-        "@xla//xla:util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/ffi:ffi_api",
         "@xla//xla/ffi/api:ffi",
diff --git a/jaxlib/xla/py_client_cpu.cc b/jaxlib/xla/py_client_cpu.cc
index fc4f895af..ac4e7bee5 100644
--- a/jaxlib/xla/py_client_cpu.cc
+++ b/jaxlib/xla/py_client_cpu.cc
@@ -41,7 +41,6 @@ limitations under the License.
 #include "xla/python/nb_numpy.h"
 #include "xla/python/types.h"
 #include "xla/shape_util.h"
-#include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
 namespace nb = nanobind;
@@ -79,9 +78,9 @@ ffi::Error XlaFfiPythonCpuCallback(FfiLoadedHostCallbacks* callbacks,
   for (size_t i = 0; i < args.size(); ++i) {
     auto arg = args.get<ffi::AnyBuffer>(i);
     auto ptype = static_cast<PrimitiveType>(arg->element_type());
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    if (ptype == S1 || ptype == U1) {
+    // TODO(b/395428868): Remove this check once we support subbyte types.
+    if (ptype == S1 || ptype == S2 || ptype == S4 || ptype == U1 ||
+        ptype == U2 || ptype == U4) {
       return ffi::Error(ffi::ErrorCode::kUnimplemented,
                         absl::StrFormat("Unsupported primitive type: %s",
                                         PrimitiveType_Name(ptype)));
@@ -97,18 +96,9 @@ ffi::Error XlaFfiPythonCpuCallback(FfiLoadedHostCallbacks* callbacks,
     auto dtype = maybe_dtype.value();
     auto dims = absl::Span<const int64_t>(arg->dimensions().begin(),
                                           arg->dimensions().size());
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    std::unique_ptr<char[]> buffer;
-    const void* data = arg->untyped_data();
-    size_t bits_per_element = xla::primitive_util::BitWidth(ptype);
-    if (bits_per_element == 2 || bits_per_element == 4) {
-      buffer = xla::UnpackIntN(bits_per_element, static_cast<const char*>(data),
-                               arg->size_bytes());
-      data = buffer.get();
-    }
     // We pass in data using default numpy layout i.e., std::nullopt.
-    auto array = nb_numpy_ndarray(dtype, dims, std::nullopt, data);
+    auto array =
+        nb_numpy_ndarray(dtype, dims, std::nullopt, arg.value().untyped_data());
     array.attr("flags").attr("writeable") = nb::bool_(false);
     PyTuple_SET_ITEM(nb_args.ptr(), i, array.release().ptr());
   }
@@ -129,9 +119,9 @@ ffi::Error XlaFfiPythonCpuCallback(FfiLoadedHostCallbacks* callbacks,
   for (size_t i = 0; i < rets.size(); ++i) {
     auto ret = rets.get<ffi::AnyBuffer>(i).value();
     auto ptype = static_cast<PrimitiveType>(ret->element_type());
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    if (ptype == S1 || ptype == U1) {
+    // TODO(b/395428868): Remove this check once we support subbyte types.
+    if (ptype == S1 || ptype == S2 || ptype == S4 || ptype == U1 ||
+        ptype == U2 || ptype == U4) {
       return ffi::Error(ffi::ErrorCode::kUnimplemented,
                         absl::StrFormat("Unsupported primitive type: %s",
                                         PrimitiveType_Name(ptype)));
@@ -151,45 +141,26 @@ ffi::Error XlaFfiPythonCpuCallback(FfiLoadedHostCallbacks* callbacks,
     }
     auto expected_shape = maybe_expected_shape.value();
     auto expected_strides = ByteStridesForShape(expected_shape);
-
-    const void* data = array.data();
-    if (strides != expected_strides) {
-      xla::TransposePlan::Options options;
-      options.elem_size_in_bytes = xla::primitive_util::ByteWidth(ptype);
-      options.dims = absl::Span<const int64_t>(
-          reinterpret_cast<const int64_t*>(array.shape()), array.ndim());
-      absl::InlinedVector<int64_t, 4> reversed_layout;
-      reversed_layout.resize(expected_shape.dimensions().size());
-      absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
-                           reversed_layout.begin());
-      options.permutation = reversed_layout;
-      options.input_layout = xla::TransposePlan::Striding{strides};
-      auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
-      if (!maybe_plan.ok()) {
-        return ffi::Error::Internal(maybe_plan.status().ToString());
-      }
-      auto plan = maybe_plan.value();
-      plan->Execute(data, ret->untyped_data());
-      data = ret->untyped_data();
-    }
-
-    // TODO(b/402422886): Remove this once we form Jax arrays directly instead
-    // of packing/unpacking to/from numpy arrays.
-    std::unique_ptr<char[]> buffer;
-    size_t bits_per_element = xla::primitive_util::BitWidth(ptype);
-    if (bits_per_element == 2 || bits_per_element == 4) {
-      // NOTE(dsuo): FFI arguments and return buffers are sized assuming
-      // minimum 1-byte element sizes, even if the data itself is packed.
-      buffer = xla::PackIntN(bits_per_element, static_cast<const char*>(data),
-                             ret->size_bytes());
-      data = buffer.get();
+    if (strides == expected_strides) {
+      std::memcpy(ret->untyped_data(), array.data(), ret->size_bytes());
+      continue;
     }
-
-    // Copy data to output buffer if haven't already or modified the data to
-    // write back.
-    if (data != ret->untyped_data()) {
-      std::memcpy(ret->untyped_data(), data, ret->size_bytes());
+    xla::TransposePlan::Options options;
+    options.elem_size_in_bytes = xla::primitive_util::ByteWidth(ptype);
+    options.dims = absl::Span<const int64_t>(
+        reinterpret_cast<const int64_t*>(array.shape()), array.ndim());
+    absl::InlinedVector<int64_t, 4> reversed_layout;
+    reversed_layout.resize(expected_shape.dimensions_size());
+    absl::c_reverse_copy(expected_shape.layout().minor_to_major(),
+                         reversed_layout.begin());
+    options.permutation = reversed_layout;
+    options.input_layout = xla::TransposePlan::Striding{strides};
+    auto maybe_plan = transpose_cache->cache.GetOrCreate(options);
+    if (!maybe_plan.ok()) {
+      return ffi::Error::Internal(maybe_plan.status().ToString());
     }
+    auto plan = maybe_plan.value();
+    plan->Execute(array.data(), ret->untyped_data());
   }
 
   return ffi::Error::Success();
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index d078359af..8b3ddfde0 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "f4a53456b04acf9b63b3b30bd828cec29c4aa7de"
-XLA_SHA256 = "2ee32b70af547fd13ce404d75c3fa9834bc8be46a488cd8f0caa10e9a6ec7ede"
+XLA_COMMIT = "8df9390dc9444d900c7c7f2c123f23b549adf8e3"
+XLA_SHA256 = "8e97c395d1e50a49fab386ccc7da1f78dc86bf670b20a892656e2e75bbf64f0e"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index 8e4ba275f..d509fe78c 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -26,7 +26,6 @@ import jax
 import jax.numpy as jnp
 from jax._src import config, test_util as jtu
 from jax._src.dtypes import _default_types, canonicalize_dtype
-from jax._src import xla_bridge as xb
 
 ARRAY_API_NAMESPACE = jnp
 
@@ -284,10 +283,7 @@ class ArrayAPIInspectionUtilsTest(jtu.JaxTestCase):
     assert self.info.default_device() is None
 
   def test_devices_info(self):
-    devices = set(self.info.devices())
-    assert None in devices
-    for backend in xb.backends():
-      assert devices.issuperset(jax.devices(backend))
+    assert self.info.devices() == jax.devices()
 
   def test_default_dtypes_info(self):
     _default_dtypes = {
diff --git a/tests/pallas/mosaic_gpu_test.py b/tests/pallas/mosaic_gpu_test.py
index 77e934f65..9c5795e49 100644
--- a/tests/pallas/mosaic_gpu_test.py
+++ b/tests/pallas/mosaic_gpu_test.py
@@ -85,13 +85,6 @@ class PallasTest(jtu.JaxTestCase, metaclass=PallasTestMetaclass):
     if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Warpgroup:
       self.skipTest("Not supported under WG semantics")
 
-  def kernel(self, *args, **kwargs):
-    compiler_params = dataclasses.replace(
-        kwargs.pop("compiler_params", plgpu.GPUCompilerParams()),
-        thread_semantics=self.THREAD_SEMANTICS,
-    )
-    return plgpu.kernel(*args, compiler_params=compiler_params, **kwargs)
-
   def pallas_call(self, *args, **kwargs):
     compiler_params = dataclasses.replace(
         kwargs.pop("compiler_params", plgpu.GPUCompilerParams()),
@@ -633,26 +626,6 @@ class PallasCallTest(PallasTest):
     x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
     np.testing.assert_array_equal(f(x), x * 2)
 
-  def test_scoped_copy_with_user_transforms(self):
-    def kernel(x_ref, o_ref, barrier_ref):
-      def body(tmp_ref):
-        tmp_ref = plgpu.unswizzle_ref(tmp_ref, 128)
-        tmp_ref = plgpu.untile_ref(tmp_ref, (8, 32))
-        plgpu.copy_gmem_to_smem(x_ref, tmp_ref, barrier_ref)
-        plgpu.barrier_wait(barrier_ref)
-        o_ref[...] = tmp_ref[...] * 2
-      pl.run_scoped(body, plgpu.SMEM((16, 4, 8, 32), jnp.float32))
-
-    in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
-    f = pl.pallas_call(
-        kernel,
-        out_shape=jax.ShapeDtypeStruct([128, 128], jnp.float32),
-        in_specs=(in_spec,),
-        scratch_shapes=[plgpu.Barrier(num_arrivals=1)],
-    )
-    x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
-    np.testing.assert_array_equal(f(x), x * 2)
-
   def test_copy_with_transforms_and_indexing(self):
     self.skip_if_wg_semantics()
 
@@ -982,7 +955,7 @@ class PallasCallTest(PallasTest):
 
   def test_load_scalar(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct((128,), jnp.int32),
         in_specs=[plgpu.GPUBlockSpec(memory_space=plgpu.GPUMemorySpace.GMEM)],
     )
@@ -994,7 +967,7 @@ class PallasCallTest(PallasTest):
 
   def test_run_scoped(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
     )
     def kernel(x_ref, o_ref):
@@ -1012,7 +985,7 @@ class PallasCallTest(PallasTest):
 
   def test_program_id(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=(),
         out_specs=pl.BlockSpec((128,), lambda *i: i),
         out_shape=jax.ShapeDtypeStruct([128 * 2], jnp.int32),
@@ -1031,7 +1004,7 @@ class PallasCallTest(PallasTest):
     # 3 CUDA grid dimensions.
     grid = (2, 3, 4, 5)
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=(),
         out_specs=pl.BlockSpec((1,) * len(grid) + (128,), lambda *i: (*i, 0)),
         out_shape=jax.ShapeDtypeStruct([*grid, 128], jnp.int32),
@@ -1052,7 +1025,7 @@ class PallasCallTest(PallasTest):
 
   def test_program_id_in_block_spec(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=(pl.BlockSpec((2, 128), lambda i: (pl.program_id(0), i)),),
         out_specs=pl.BlockSpec((2, 128), lambda i: (pl.program_id(0), i)),
         out_shape=jax.ShapeDtypeStruct([2, 128], jnp.int32),
@@ -1066,7 +1039,7 @@ class PallasCallTest(PallasTest):
 
   def test_num_programs(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=(),
         out_specs=pl.BlockSpec((128,), lambda *i: i),
         out_shape=jax.ShapeDtypeStruct([128 * 2], jnp.int32),
@@ -1081,7 +1054,6 @@ class PallasCallTest(PallasTest):
     )
 
   def test_swizzled_blockspec_shapes(self):
-    self.skip_if_wg_semantics()
 
     spec = plgpu.GPUBlockSpec(
         (128, 64),
@@ -1091,7 +1063,7 @@ class PallasCallTest(PallasTest):
         ),
     )
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=[spec],
         out_specs=spec,
         out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float16),
@@ -1132,7 +1104,7 @@ class PallasCallTest(PallasTest):
 
   def test_fori_loop_dynamic_bounds(self):
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
         grid=(1,)
     )
@@ -1209,10 +1181,8 @@ class PallasCallTest(PallasTest):
     )
 
   def test_while_loop_layout_mismatch(self):
-    self.skip_if_wg_semantics()  # `plgpu.layout_cast` is not supported.
-
     @functools.partial(
-        self.pallas_call, out_shape=jax.ShapeDtypeStruct([128], jnp.int32)
+        pl.pallas_call, out_shape=jax.ShapeDtypeStruct([128], jnp.int32)
     )
     def kernel(o_ref):
       def cond(acc):
@@ -1265,10 +1235,8 @@ class PallasCallTest(PallasTest):
     x = jnp.arange(256, dtype=jnp.int32)
     np.testing.assert_array_equal(kernel(x), jnp.broadcast_to(jnp.sum(x) * 3, [256]))
 
+  # Not testing with warpgroup semantics, because we want to enforce a layout.
   def test_tile_slicing(self):
-    # Not testing with warpgroup semantics, because we want to enforce a layout.
-    self.skip_if_wg_semantics()
-
     shape = (256, 128)
     block_spec = plgpu.GPUBlockSpec(
         transforms=(
@@ -1276,7 +1244,7 @@ class PallasCallTest(PallasTest):
         )
     )
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         in_specs=[block_spec],
         out_specs=block_spec,
         out_shape=jax.ShapeDtypeStruct((64, 64), jnp.uint16),
@@ -1301,7 +1269,7 @@ class PallasCallTest(PallasTest):
       a_ref[...] = jnp.ones_like(a_ref)
 
     a = np.zeros((64, 64), dtype=jnp.float32)
-    b = self.pallas_call(
+    b = pl.pallas_call(
         kernel,
         in_specs=[plgpu.GPUBlockSpec(memory_space=plgpu.GPUMemorySpace.GMEM)],
         out_specs=plgpu.GPUBlockSpec(memory_space=plgpu.GPUMemorySpace.GMEM),
@@ -1311,8 +1279,6 @@ class PallasCallTest(PallasTest):
     np.testing.assert_array_equal(b, np.ones_like(a))
 
   def test_slicing(self):
-    self.skip_if_wg_semantics()
-
     left = upper = slice(None, 64)
     right = lower = slice(64, None)
     # We rotate the four quadrants of the input clockwise.
@@ -1331,16 +1297,14 @@ class PallasCallTest(PallasTest):
             plgpu.SwizzleTransform(128),
         ),
     )
-    f = self.pallas_call(rotate, out_shape=x, in_specs=[spec], out_specs=spec)
+    f = pl.pallas_call(rotate, out_shape=x, in_specs=[spec], out_specs=spec)
     expected = np.empty_like(x)
     rotate(x, expected)
     np.testing.assert_array_equal(f(x), expected)
 
   def test_layout_cast(self, shape=(256, 64)):
-    self.skip_if_wg_semantics()  # `plgpu.layout_cast` is not supported.
-
     @functools.partial(
-        self.pallas_call,
+        pl.pallas_call,
         out_shape=jax.ShapeDtypeStruct(shape, jnp.float32),
     )
     def kernel(o_ref):
@@ -1350,8 +1314,6 @@ class PallasCallTest(PallasTest):
     np.testing.assert_array_equal(kernel(), x)
 
   def test_profiler(self):
-    self.skip_if_wg_semantics()  # Transform inference fails.
-
     def kernel(x_ref, o_ref):
       with jax.named_scope("add"):
         with jax.named_scope("load"):
@@ -1361,7 +1323,7 @@ class PallasCallTest(PallasTest):
         o_ref[...] = o
     with tempfile.TemporaryDirectory() as tmpdir:
       x = jnp.arange(256).astype(jnp.float32)
-      y = self.pallas_call(
+      y = pl.pallas_call(
           kernel,
           out_shape=jax.ShapeDtypeStruct([256], jnp.float32),
           compiler_params=plgpu.GPUCompilerParams(
@@ -1405,6 +1367,9 @@ class PallasCallTest(PallasTest):
     )
 
   def test_optimization_barrier(self):
+    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
+      self.skipTest("This test crashes with lane semantics")
+
     @functools.partial(
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
@@ -1416,6 +1381,9 @@ class PallasCallTest(PallasTest):
     np.testing.assert_array_equal(kernel(x), x)
 
   def test_optimization_barrier_multiple_inputs(self):
+    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
+      self.skipTest("This test crashes with lane semantics")
+
     @functools.partial(
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct((128,), jnp.float32),
@@ -1459,15 +1427,10 @@ class PallasCallSm90ATest(PallasSm90ATest):
 
   @parameterized.parameters(False, True)
   def test_fori_loop_accumulator(self, force_while):
-    # ``pl.run_state`` is not supported in WG semantics.
-    self.skip_if_wg_semantics()
-
     transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
     @functools.partial(
-        self.pallas_call,
-        in_specs=[
-            plgpu.GPUBlockSpec((64, 64), lambda: (0, 0), transforms=transforms)
-        ],
+        pl.pallas_call,
+        in_specs=[plgpu.GPUBlockSpec((64, 64), lambda: (0, 0), transforms=transforms)],
         out_shape=jax.ShapeDtypeStruct((64, 64), jnp.float16),
         out_specs=plgpu.GPUBlockSpec((64, 64), lambda: (0, 0)),
     )
@@ -1479,7 +1442,8 @@ class PallasCallSm90ATest(PallasSm90ATest):
     acc_ini = jnp.ones((64, 64), dtype=jnp.float16)
     np.testing.assert_array_equal(kernel(acc_ini), jnp.full((64, 64), 5, dtype=jnp.float16))
 
-  def test_realistic_matmul(self):
+  @parameterized.parameters([*plgpu.ThreadSemantics])
+  def test_realistic_matmul(self, thread_semantics):
     dtype = jnp.float16
     swizzle = 128
     elems_128b = swizzle // jnp.dtype(dtype).itemsize
@@ -1516,7 +1480,7 @@ class PallasCallSm90ATest(PallasSm90ATest):
         lambda m, n, k: (m, n),
     )
 
-    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
+    if thread_semantics == plgpu.ThreadSemantics.Lane:
       lhs_spec = plgpu.GPUBlockSpec(
           lhs_spec.block_shape, lhs_spec.index_map,
           transforms=(
@@ -1539,7 +1503,7 @@ class PallasCallSm90ATest(PallasSm90ATest):
           )
       )
 
-    res = self.pallas_call(
+    res = pl.pallas_call(
         kernel,
         in_specs=[lhs_spec, rhs_spec],
         out_specs=out_spec,
@@ -1550,14 +1514,13 @@ class PallasCallSm90ATest(PallasSm90ATest):
             dimension_semantics=["parallel", "parallel", "sequential"],
             max_concurrent_steps=2,
             delay_release=1,
+            thread_semantics=thread_semantics,
         ),
     )(a, b)
     np.testing.assert_allclose(res, a @ b, rtol=1e-3)
 
   @parameterized.parameters(jnp.float16, jnp.float32)
   def test_wgmma(self, dtype):
-    self.skip_if_wg_semantics()
-
     # TensorCores can only fuse transposes of 16-bit values, and RHS
     # is expected to be column major by default.
     rhs_transpose = jnp.dtype(dtype).itemsize != 2
@@ -1580,7 +1543,9 @@ class PallasCallSm90ATest(PallasSm90ATest):
     b = jax.random.uniform(key2, shape=b_shape, dtype=dtype)
 
     rhs_transforms = (plgpu.TilingTransform((8, elems_128b)),)
-    res = self.pallas_call(
+    if rhs_transpose:
+      rhs_transforms += (plgpu.TransposeTransform((1, 0, 2, 3)),)
+    res = pl.pallas_call(
         kernel,
         in_specs=[
             plgpu.GPUBlockSpec(
@@ -1616,18 +1581,12 @@ class PallasCallSm90ATest(PallasSm90ATest):
     a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
     b = jax.random.uniform(key2, shape=(128, 192), dtype=jnp.float16)
 
-    transforms = ()
-    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
-      transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
-    res = self.pallas_call(
+    transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
+    res = pl.pallas_call(
         kernel,
         in_specs=[
-            plgpu.GPUBlockSpec(
-                (64, 128), lambda: (0, 0), transforms=transforms
-            ),
-            plgpu.GPUBlockSpec(
-                (128, 192), lambda: (0, 0), transforms=transforms
-            ),
+            plgpu.GPUBlockSpec((64, 128), lambda: (0, 0), transforms=transforms),
+            plgpu.GPUBlockSpec((128, 192), lambda: (0, 0), transforms=transforms),
         ],
         out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
         out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float32),
@@ -1635,9 +1594,6 @@ class PallasCallSm90ATest(PallasSm90ATest):
     np.testing.assert_allclose(res, a @ b, rtol=1e-3)
 
   def test_wgmma_registers_init(self):
-    # ``pl.run_state`` is not supported in WG semantics.
-    self.skip_if_wg_semantics()
-
     def kernel(a_ref, b_ref, i_ref, o_ref):
       def scope(acc_ref):
         plgpu.wgmma(acc_ref, a_ref[...], b_ref)
@@ -1649,18 +1605,12 @@ class PallasCallSm90ATest(PallasSm90ATest):
     i = jax.random.uniform(key3, shape=(64, 192), dtype=jnp.float16) * 10
 
     transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
-    res = self.pallas_call(
+    res = pl.pallas_call(
         kernel,
         in_specs=[
-            plgpu.GPUBlockSpec(
-                (64, 128), lambda: (0, 0), transforms=transforms
-            ),
-            plgpu.GPUBlockSpec(
-                (128, 192), lambda: (0, 0), transforms=transforms
-            ),
-            plgpu.GPUBlockSpec(
-                (64, 192), lambda: (0, 0), transforms=transforms
-            ),
+            plgpu.GPUBlockSpec((64, 128), lambda: (0, 0), transforms=transforms),
+            plgpu.GPUBlockSpec((128, 192), lambda: (0, 0), transforms=transforms),
+            plgpu.GPUBlockSpec((64, 192), lambda: (0, 0), transforms=transforms),
         ],
         out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
         out_shape=jax.ShapeDtypeStruct((64, 192), jnp.float16),
@@ -1668,8 +1618,6 @@ class PallasCallSm90ATest(PallasSm90ATest):
     np.testing.assert_allclose(res, i + a @ b, rtol=2e-3)
 
   def test_wgmma_sliced_ref(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     def kernel(a_ref, b_ref, o_ref):
       def scope(acc_ref):
         plgpu.wgmma(acc_ref, a_ref.at[0], b_ref.at[0])
@@ -1681,18 +1629,22 @@ class PallasCallSm90ATest(PallasSm90ATest):
     a = jax.random.uniform(key1, shape=(2, 64, 128), dtype=jnp.float16)
     b = jax.random.uniform(key2, shape=(2, 128, 192), dtype=jnp.float16)
 
-    transforms = ()
-    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
-      transforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))
-
-    res = self.pallas_call(
+    res = pl.pallas_call(
         kernel,
         in_specs=[
             plgpu.GPUBlockSpec(
-                (2, 64, 128), lambda: (0, 0, 0), transforms=transforms
+                (2, 64, 128), lambda: (0, 0, 0),
+                transforms=(
+                    plgpu.TilingTransform((8, 64)),
+                    plgpu.SwizzleTransform(128),
+                ),
             ),
             plgpu.GPUBlockSpec(
-                (2, 128, 192), lambda: (0, 0, 0), transforms=transforms
+                (2, 128, 192), lambda: (0, 0, 0),
+                transforms=(
+                    plgpu.TilingTransform((8, 64)),
+                    plgpu.SwizzleTransform(128),
+                ),
             ),
         ],
         out_specs=plgpu.GPUBlockSpec((64, 192), lambda: (0, 0)),
@@ -1701,8 +1653,6 @@ class PallasCallSm90ATest(PallasSm90ATest):
     np.testing.assert_allclose(res, a[0] @ b[0], rtol=1e-3)
 
   def test_wgmma_sliced_acc(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     swizzle = 128
     elems_128b = swizzle // jnp.dtype(jnp.float16).itemsize
     def kernel(a_ref, b_ref, o_ref):
@@ -1715,41 +1665,38 @@ class PallasCallSm90ATest(PallasSm90ATest):
     key1, key2 = jax.random.split(jax.random.key(42), 2)
     a = jax.random.uniform(key1, shape=(64, 128), dtype=jnp.float16)
     b = jax.random.uniform(key2, shape=(128, 128), dtype=jnp.float16)
-    transforms = ()
-    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
-      transforms = (
-          plgpu.TilingTransform((8, elems_128b)),
-          plgpu.SwizzleTransform(128),
-      )
-    res = self.pallas_call(
+    res = pl.pallas_call(
         kernel,
         in_specs=[
             plgpu.GPUBlockSpec(
-                (64, 128), lambda *ij: ij, transforms=transforms
+                (64, 128),
+                lambda i, j: (i, j),
+                transforms=(
+                    plgpu.TilingTransform((8, elems_128b)),
+                    plgpu.SwizzleTransform(128),
+                ),
             ),
             plgpu.GPUBlockSpec(
-                (128, 128), lambda *ij: ij, transforms=transforms
+                (128, 128),
+                lambda *i: i,
+                transforms=(
+                    plgpu.TilingTransform((8, elems_128b)),
+                    plgpu.SwizzleTransform(128),
+                ),
             ),
         ],
-        out_specs=plgpu.GPUBlockSpec((64, 128), lambda *ij: ij),
+        out_specs=plgpu.GPUBlockSpec((64, 128), lambda *i: i),
         out_shape=jax.ShapeDtypeStruct((64, 128), jnp.float32),
         grid=(1, 1),
     )(a, b)
     np.testing.assert_allclose(res, a @ b, rtol=1e-3)
 
 
-class PallasCallSm90AWGTest(
-    PallasCallSm90ATest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class PallasCallSm100ATest(PallasSm100ATest):
 
   def test_tmem_alloc(self):
-
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros((128, 128), jnp.float32),
         scratch_shapes=[
             plgpu.TMEM((128, 128), jnp.float32),
@@ -1769,12 +1716,6 @@ class PallasCallSm100ATest(PallasSm100ATest):
     jax.block_until_ready(kernel())
 
 
-class PallasCallSm100AWGTest(
-    PallasCallSm100ATest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class PipelineTest(PallasTest):
 
   def test_pipeline_mode(self):
@@ -1796,13 +1737,13 @@ class PipelineTest(PallasTest):
 
     @jax.jit
     def vadd(x, y):
-      return self.pallas_call(
-          body,
-          out_shape=jax.ShapeDtypeStruct(x.shape, jnp.float32),
-          in_specs=in_specs,
-          out_specs=out_specs,
-          grid=data_size // block_size,
-      )(x, y)
+      return pl.pallas_call(
+        body,
+        out_shape=jax.ShapeDtypeStruct(x.shape, jnp.float32),
+        in_specs=in_specs,
+        out_specs=out_specs,
+        grid=data_size // block_size,
+    )(x, y)
 
     with self.assertRaisesRegex(Exception, "Pipeline mode is not supported"):
       vadd(x, y)
@@ -1864,7 +1805,7 @@ class PipelineTest(PallasTest):
       plgpu.wait_smem_to_gmem(0)
 
     x = jnp.arange(32 * 4 * 64).reshape(32 * 4, 64).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -1878,9 +1819,6 @@ class PipelineTest(PallasTest):
       ((plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(128)),),
   )
   def test_emit(self, transforms):
-    if transforms:
-      self.skip_if_wg_semantics()
-
     num_steps = 4
 
     def kernel(x_gmem, o_gmem):
@@ -1907,7 +1845,7 @@ class PipelineTest(PallasTest):
 
     x = jnp.arange(64 * num_steps * 64)
     x = x.reshape(-1, num_steps * 64).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -1940,7 +1878,7 @@ class PipelineTest(PallasTest):
 
     x = jnp.arange(32 * num_steps * 16)
     x = x.reshape(-1, num_steps * 16).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -1965,7 +1903,7 @@ class PipelineTest(PallasTest):
 
     x = jnp.arange(32 * num_steps * 16)
     x = x.reshape(-1, num_steps * 16).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -1998,7 +1936,7 @@ class PipelineTest(PallasTest):
 
     x = jnp.arange(num_steps1 * 32 * num_steps2 * 16)
     x = x.reshape(-1, num_steps2 * 16).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -2026,7 +1964,7 @@ class PipelineTest(PallasTest):
 
     x = jnp.arange(32 * num_steps1 * 16 * num_steps2 * 8)
     x = x.reshape(-1, num_steps1 * 16, num_steps2 * 8).astype(jnp.float32)
-    kernel_fn = self.pallas_call(
+    kernel_fn = pl.pallas_call(
         kernel,
         in_specs=[pl.BlockSpec(memory_space=plgpu.GMEM)],
         out_specs=pl.BlockSpec(memory_space=plgpu.GMEM),
@@ -2035,17 +1973,9 @@ class PipelineTest(PallasTest):
     np.testing.assert_array_equal(kernel_fn(x), x + 1.0)
 
 
-class PipelineWGTest(
-    PipelineTest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class PipelineSm90ATest(PallasSm90ATest):
 
   def test_realistic_matmul(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     dtype = jnp.float16
     swizzle = 128
     elems_128b = swizzle // jnp.dtype(dtype).itemsize
@@ -2055,13 +1985,6 @@ class PipelineSm90ATest(PallasSm90ATest):
     tile_k = elems_128b
     m, k, n = grid_m * tile_m, grid_k * tile_k, grid_n * tile_n
 
-    transforms = ()
-    if self.THREAD_SEMANTICS == plgpu.ThreadSemantics.Lane:
-      transforms = (
-          plgpu.TilingTransform((8, elems_128b)),
-          plgpu.SwizzleTransform(128),
-      )
-
     def kernel(a_gmem, b_gmem, o_smem, acc):
       def kernel_body(a_smem, b_smem):
         assert a_smem.shape == (tile_m, tile_k)
@@ -2075,11 +1998,21 @@ class PipelineSm90ATest(PallasSm90ATest):
           kernel_body,
           in_specs=[
               plgpu.GPUBlockSpec(
-                  (tile_m, tile_k), lambda k: (pid_m, k), transforms=transforms
-              ),
-              plgpu.GPUBlockSpec(
-                  (tile_k, tile_n), lambda k: (k, pid_n), transforms=transforms
-              ),
+                (tile_m, tile_k),
+                lambda k: (pid_m, k),
+                transforms=(
+                    plgpu.TilingTransform((8, elems_128b)),
+                    plgpu.SwizzleTransform(128),
+                ),
+            ),
+            plgpu.GPUBlockSpec(
+                (tile_k, tile_n),
+                lambda k: (k, pid_n),
+                transforms=(
+                    plgpu.TilingTransform((8, elems_128b)),
+                    plgpu.SwizzleTransform(128),
+                ),
+            ),
           ],
           grid=(grid_k,),
           max_concurrent_steps=2,
@@ -2092,14 +2025,19 @@ class PipelineSm90ATest(PallasSm90ATest):
     a = jax.random.uniform(key1, shape=(m, k), dtype=dtype)
     b = jax.random.uniform(key2, shape=(k, n), dtype=dtype)
 
-    res = self.pallas_call(
+    res = pl.pallas_call(
         kernel,
         in_specs=[
             pl.BlockSpec(memory_space=plgpu.GMEM),
-            pl.BlockSpec(memory_space=plgpu.GMEM),
+            pl.BlockSpec(memory_space=plgpu.GMEM)
         ],
         out_specs=plgpu.GPUBlockSpec(
-            (tile_m, tile_n), lambda m, n: (m, n), transforms=transforms
+            (tile_m, tile_n),
+            lambda m, n: (m, n),
+            transforms=(
+                plgpu.TilingTransform((8, elems_128b)),
+                plgpu.SwizzleTransform(128),
+            ),
         ),
         out_shape=jax.ShapeDtypeStruct((m, n), jnp.float16),
         scratch_shapes=[plgpu.ACC((tile_m, tile_n), jnp.float32)],
@@ -2108,19 +2046,11 @@ class PipelineSm90ATest(PallasSm90ATest):
     np.testing.assert_array_equal(res, a @ b)
 
 
-class PipelineSm90AWGTest(
-    PipelineSm90ATest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class WarpSpecializedPipelineTest(PallasTest):
 
   @parameterized.product(m=[512], n=[512],
                          manual_consumed_barriers=[False, True])
   def test_pipelined_copy(self, m, n, manual_consumed_barriers):
-    self.skip_if_wg_semantics()  # Times out!
-
     x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float16)
     blk_m = blk_n = 64
 
@@ -2154,7 +2084,7 @@ class WarpSpecializedPipelineTest(PallasTest):
             ),
         ],
     )
-    kernel = self.kernel(
+    kernel = plgpu.kernel(
         pipeline,
         out_shape=(
             jax.ShapeDtypeStruct((m, n), jnp.float16),
@@ -2171,8 +2101,6 @@ class WarpSpecializedPipelineTest(PallasTest):
     np.testing.assert_array_equal(out_last_block, x[-blk_m:, -blk_n:])
 
   def test_elementwise_add(self, m=256, n=256, num_compute_wgs=2):
-    self.skip_if_wg_semantics()  # Crashes!
-
     blk_m = blk_n = 64
     spec = pl.BlockSpec(
         block_shape=(blk_m, blk_n), index_map=lambda i, j: (i, j)
@@ -2194,7 +2122,7 @@ class WarpSpecializedPipelineTest(PallasTest):
         in_specs=[spec, spec],
         out_specs=[spec],
     )
-    kernel = self.kernel(
+    kernel = plgpu.kernel(
         pipeline,
         out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
         compiler_params=plgpu.GPUCompilerParams(approx_math=True),
@@ -2208,12 +2136,10 @@ class WarpSpecializedPipelineTest(PallasTest):
     np.testing.assert_allclose(kernel(x, y), x + y, atol=1e-4)
 
   def test_carry_accumulate(self, m=256, n=256, num_compute_wgs=2):
-    self.skip_if_wg_semantics()  # `plgpu.layout_cast` is not supported.
-
     blk_m = blk_n = 64
 
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jax.ShapeDtypeStruct((blk_m, blk_n), jnp.float32),
         scratch_shapes=[
             plgpu.SMEM((blk_m, blk_n), jnp.float32),
@@ -2270,19 +2196,11 @@ class WarpSpecializedPipelineTest(PallasTest):
     np.testing.assert_allclose(kernel(x), ref, atol=1e-4)
 
 
-class WarpSpecializedPipelineWGTest(
-    WarpSpecializedPipelineTest,
-    thread_semantics=plgpu.ThreadSemantics.Warpgroup,
-):
-  ...
-
-
 class CoreMapTest(PallasTest):
 
   def test_multiple_wg(self):
-
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros((2, 128), np.int32),
         num_threads=2,
         thread_name="wg",
@@ -2296,9 +2214,8 @@ class CoreMapTest(PallasTest):
     )
 
   def test_multiple_wg_with_grid(self):
-
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros((4, 2, 128), np.int32),
         grid=(2, 2),
         grid_names=("x", "y"),
@@ -2328,7 +2245,7 @@ class CoreMapTest(PallasTest):
     num_threads = 2
 
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros(
             (b, x_dim, y_dim, z_dim, num_threads, 128), np.int32
         ),
@@ -2355,10 +2272,8 @@ class CoreMapTest(PallasTest):
     np.testing.assert_array_equal(result, ref)
 
   def test_cross_wg_barrier(self):
-    self.skip_if_wg_semantics()  # Times out!
-
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros((2, 128), np.int32),
         # Each warpgroup is a single logical thread!
         scratch_shapes=[plgpu.Barrier(num_arrivals=2)],
@@ -2376,10 +2291,8 @@ class CoreMapTest(PallasTest):
     )
 
   def test_cluster(self):
-    self.skip_if_wg_semantics()  # Needs debug_print in the MGPU dialect.
-
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jnp.zeros(128, np.int32),
         grid=(2,),
         grid_names=("x",),
@@ -2406,8 +2319,6 @@ class CoreMapTest(PallasTest):
     )
 
   def test_realistic_matmul_with_cluster(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     dtype = jnp.float16
     swizzle = 128
     elems_128b = swizzle // jnp.dtype(dtype).itemsize
@@ -2432,7 +2343,7 @@ class CoreMapTest(PallasTest):
     delay_release = 1
 
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jax.ShapeDtypeStruct((m, n), dtype),
         scratch_shapes=[
             plgpu.SMEM(
@@ -2529,44 +2440,34 @@ class CoreMapTest(PallasTest):
     np.testing.assert_array_equal(kernel(a, b), a @ b)
 
 
-class CoreMapWGTest(
-    CoreMapTest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class ExamplesTest(PallasTest):
 
   # Basic
   def test_stage0(self):
-    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
-
-    @functools.partial(self.kernel, out_shape=x)
-    def kernel(l_ref, r_ref, o_ref):
+    def body(l_ref, r_ref, o_ref):
       o_ref[...] = l_ref[...] + r_ref[...]
 
-    np.testing.assert_allclose(kernel(x, x), x + x)
+    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
+    out = plgpu.kernel(body, out_shape=x)(x, x)
+    np.testing.assert_allclose(out, x + x)
 
   # Multi-block kernels
   def test_stage1(self):
     row_block = 64
-    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
-
-    @functools.partial(
-        self.kernel, out_shape=x, grid=(2,), grid_names=("rows",)
-    )
-    def kernel(l_ref, r_ref, o_ref):
+    def body(l_ref, r_ref, o_ref):
       my_slice = pl.ds(lax.axis_index("rows") * row_block, row_block)
       o_ref[my_slice] = l_ref[my_slice] + r_ref[my_slice]
 
-    np.testing.assert_allclose(kernel(x, x), x + x)
+    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
+    out = plgpu.kernel(body, out_shape=x, grid=(2,), grid_names=("rows",))(x, x)
+    np.testing.assert_allclose(out, x + x)
 
   # Async copies
   def test_stage3(self):
     row_block, col_block = 64, 128
 
     @functools.partial(
-        self.kernel,
+        plgpu.kernel,
         out_shape=jax.ShapeDtypeStruct((128, 128), jnp.float16),
         scratch_shapes=[
             *([plgpu.SMEM((row_block, col_block), jnp.float16)] * 3),
@@ -2591,12 +2492,7 @@ class ExamplesTest(PallasTest):
   # Pipelining
   def test_stage4(self):
     row_block, col_block = 64, 32
-    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
-
-    @functools.partial(
-        self.kernel, out_shape=x, grid=(2,), grid_names=("rows",)
-    )
-    def kernel(l_ref, r_ref, o_ref):
+    def body(l_ref, r_ref, o_ref):
       def compute(l_smem, r_smem, o_smem):
         o_smem[...] = l_smem[...] + r_smem[...]
       r = lax.axis_index("rows")
@@ -2608,19 +2504,14 @@ class ExamplesTest(PallasTest):
           out_specs=[block],
       )(l_ref, r_ref, o_ref)
 
-    np.testing.assert_allclose(kernel(x, x), x + x)
+    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
+    out = plgpu.kernel(body, out_shape=x, grid=(2,), grid_names=("rows",))(x, x)
+    np.testing.assert_allclose(out, x + x)
 
   # Transforms
   def test_stage5(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     row_block, col_block = 64, 32
-    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
-
-    @functools.partial(
-        self.kernel, out_shape=x, grid=(2,), grid_names=("rows",)
-    )
-    def kernel(l_ref, r_ref, o_ref):
+    def body(l_ref, r_ref, o_ref):
       def compute(l_smem, r_smem, o_smem):
         o_smem[...] = l_smem[...] + r_smem[...]
       r = lax.axis_index("rows")
@@ -2635,7 +2526,9 @@ class ExamplesTest(PallasTest):
           out_specs=[block],
       )(l_ref, r_ref, o_ref)
 
-    np.testing.assert_allclose(kernel(x, x), x + x)
+    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
+    out = plgpu.kernel(body, out_shape=x, grid=(2,), grid_names=("rows",))(x, x)
+    np.testing.assert_allclose(out, x + x)
 
   def test_semaphore_lowering(self):
     # This is a smoke test until we add support for lowering of semaphore ops.
@@ -2645,10 +2538,8 @@ class ExamplesTest(PallasTest):
       assert jnp.issubdtype(sem_ref.dtype, pl.semaphore)
       o_ref[...] = i_ref1[...]
     x = jnp.arange(128, dtype=jnp.float32).reshape((128,))
-    kernel = self.pallas_call(
-        body,
-        out_shape=x,
-        scratch_shapes=[plgpu.SemaphoreType.REGULAR((4,))],
+    kernel = pl.pallas_call(
+        body, out_shape=x, scratch_shapes=[plgpu.SemaphoreType.REGULAR((4,))],
     )
     text = jax.jit(kernel).lower(x, x).as_text()
     self.assertIn(
@@ -2664,33 +2555,19 @@ class ExamplesTest(PallasTest):
     )
 
 
-class ExamplesWGTest(
-    ExamplesTest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 class ExamplesSm90ATest(PallasSm90ATest):
 
   # WGMMA
   def test_stage6(self):
-    self.skip_if_wg_semantics()  # Needs WGMMA to support slices.
-
     m_block = n_block = 64
     k_block = 32
-    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
-
-    @functools.partial(
-        self.kernel, out_shape=x, grid=(2, 2), grid_names=("m", "n")
-    )
-    def kernel(l_ref, r_ref, o_ref):
+    def body(l_ref, r_ref, o_ref):
       def compute(l_smem, r_smem, o_smem):
         def do_wgmma(acc_ref):
           plgpu.wgmma(acc_ref, l_smem, r_smem)
           return acc_ref[...]
         o_smem[...] += pl.run_scoped(do_wgmma, plgpu.ACC((m_block, n_block), jnp.float16))
-      m = lax.axis_index("m")
-      n = lax.axis_index("n")
+      m, n = lax.axis_index("m"), lax.axis_index("n")
       lo_transforms = (plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(64))
       r_transforms = (plgpu.TilingTransform((8, 32)), plgpu.SwizzleTransform(64))
       plgpu.emit_pipeline(
@@ -2701,16 +2578,12 @@ class ExamplesSm90ATest(PallasSm90ATest):
           out_specs=[plgpu.GPUBlockSpec((m_block, n_block), lambda k: (m, n), transforms=lo_transforms)],
       )(l_ref, r_ref, o_ref)
 
-    np.testing.assert_allclose(kernel(x, x), x @ x)
+    x = jnp.arange(128 * 128, dtype=jnp.float16).reshape(128, 128)
+    out = plgpu.kernel(body, out_shape=x, grid=(2, 2), grid_names=("m", "n"))(x, x)
+    np.testing.assert_allclose(out, x @ x)
 
   # TODO(apaszke): Clusters and multicast
 
 
-class ExamplesSm90AWGTest(
-    ExamplesSm90ATest, thread_semantics=plgpu.ThreadSemantics.Warpgroup
-):
-  ...
-
-
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/pgle_test.py b/tests/pgle_test.py
index 7dabd809d..2787de4c6 100644
--- a/tests/pgle_test.py
+++ b/tests/pgle_test.py
@@ -21,7 +21,7 @@ import shutil
 import tempfile
 import warnings
 
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 import jax
 from jax._src import api
 from jax._src import compilation_cache as cc
@@ -478,5 +478,55 @@ class PgleTest(jtu.JaxTestCase):
           self.assertLen(w, 1)
           self.assertIn("PERSISTENT CACHE WRITE with key jit_h-", str(w[0].message))
 
+  @parameterized.parameters([True, False])
+  @jtu.thread_unsafe_test()
+  def testAutoPgleWithCommandBuffers(self, enable_compilation_cache):
+    with (config.pgle_profiling_runs(1),
+          config.enable_compilation_cache(enable_compilation_cache),
+          config.enable_pgle(True),
+          tempfile.TemporaryDirectory() as dump_dir,
+          tempfile.TemporaryDirectory() as cache_dir):
+      if enable_compilation_cache:
+        cc.reset_cache()
+        cc.set_cache_dir(cache_dir)
+      compiler_options = {
+        'xla_dump_to': dump_dir,
+        # FUSION, see https://github.com/openxla/xla/issues/22459
+        'xla_gpu_enable_command_buffer': 1,
+        'xla_gpu_graph_min_graph_size': 1,
+      }
+      @partial(
+          jax.jit,
+          compiler_options=compiler_options,
+      )
+      def f(x):
+        return x * 2
+
+      x = jnp.arange(1)
+      expected = x * 2
+
+      # This is ugly, but it does not seem possible to get the AutoPGLE-recompiled
+      # executable text (.lower(x).compile().as_text() or similar).
+      def get_new_hlo():
+        additions = set(os.listdir(dump_dir)) - get_new_hlo.seen_files
+        get_new_hlo.seen_files |= additions
+        new_hlos = list(filter(lambda f: f.endswith("_gpu_after_optimizations.txt"), additions))
+        assert len(new_hlos) == 1
+        with open(os.path.join(dump_dir, new_hlos[0]), "r") as ifile:
+          return ifile.read()
+
+      get_new_hlo.seen_files = set()
+
+      # Run 1
+      self.assertArraysEqual(f(x), expected)
+      self.assertNotIn("command_buffer", get_new_hlo()) # b/376647494 workaround
+      # Run 2
+      self.assertArraysEqual(f(x), expected)
+      self.assertIn("command_buffer", get_new_hlo()) # workaround disabled
+
+    api.clear_caches()
+    pjit._pgle_profiler_dict.clear()
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index 34ab20c05..a8442b4a1 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -586,15 +586,10 @@ class PythonCallbackTest(jtu.JaxTestCase):
         self.assertAllClose(2 * x, fun(x))
     self.assertEqual(count(), 1)
 
-  @parameterized.parameters("int2", "int4", "uint2", "uint4", "float4_e2m1fn")
+  @parameterized.parameters("int2", "int4", "uint2", "uint4")
   def test_subbyte_operands(self, dtype: str):
-    if jaxlib_extension_version < 323:
-      self.skipTest("Requires jaxlib_extension_version >= 323.")
-    if "2" in dtype and jtu.test_device_matches(["tpu"]):
-      self.skipTest(
-          "TODO(dsuo): TPU callbacks send SIGABRT for int2, uint2, and"
-          " float4_e2m1fn."
-      )
+    if jaxlib_extension_version <= 321:
+      self.skipTest("Requires jaxlib_extension_version >= 322.")
     def get(x):
       return x
     def f(x):
@@ -605,17 +600,21 @@ class PythonCallbackTest(jtu.JaxTestCase):
       )
       return y
     x = np.arange(8, dtype=dtype)
-    np.testing.assert_array_equal(jax.jit(f)(x), np.arange(8, dtype=dtype))
+    # TODO(b/395428868): Remove this check once we support subbyte types.
+    if jtu.test_device_matches(["tpu"]):
+      if "2" in dtype:
+        self.skipTest("TODO(dsuo): TPU callbacks send SIGABRT for int2/uint2.")
+      np.testing.assert_array_equal(jax.jit(f)(x), np.arange(8, dtype=dtype))
+    else:
+      with self.assertRaisesRegex(
+          Exception, "Unsupported primitive type"
+      ):
+        _ = jax.jit(f)(x)
 
-  @parameterized.parameters("int2", "int4", "uint2", "uint4", "float4_e2m1fn")
+  @parameterized.parameters("int2", "int4", "uint2", "uint4")
   def test_subbyte_results(self, dtype: str):
-    if jaxlib_extension_version < 323:
-      self.skipTest("Requires jaxlib_extension_version >= 323.")
-    if "2" in dtype and jtu.test_device_matches(["tpu"]):
-      self.skipTest(
-          "TODO(dsuo): TPU callbacks send SIGABRT for int2, uint2, and"
-          " float4_e2m1fn."
-      )
+    if jaxlib_extension_version <= 321:
+      self.skipTest("Requires jaxlib_extension_version >= 322.")
     def get():
       return np.arange(8, dtype=dtype)
 
@@ -626,43 +625,16 @@ class PythonCallbackTest(jtu.JaxTestCase):
       )
       return y
 
-    np.testing.assert_array_equal(jax.jit(f)(), np.arange(8, dtype=dtype))
-
-  @parameterized.parameters("int2", "int4", "uint2", "uint4", "float4_e2m1fn")
-  def test_non_default_stride_subbyte_results(self, dtype: str):
-    if jaxlib_extension_version < 323:
-      self.skipTest("Requires jaxlib_extension_version >= 323.")
-    if "2" in dtype and jtu.test_device_matches(["tpu"]):
-      self.skipTest(
-          "TODO(dsuo): TPU callbacks send SIGABRT for int2, uint2, and"
-          " float4_e2m1fn."
-      )
-    x = jnp.arange(24, dtype=dtype).reshape(2, 3, 4)
-    def callback(x):
-      return np.asfortranarray(x)
-
-    @jax.jit
-    def f(x):
-      return jax.pure_callback(
-          callback, jax.ShapeDtypeStruct(x.shape, x.dtype), x
-      )
-
-    result = f(x)
-    np.testing.assert_array_equal(x, result)
-
-  def test_non_default_stride(self):
-    x = jnp.arange(24, dtype=jnp.float32).reshape(2, 3, 4)
-    def callback(x):
-      return np.asfortranarray(x)
-
-    @jax.jit
-    def f(x):
-      return jax.pure_callback(
-          callback, jax.ShapeDtypeStruct(x.shape, x.dtype), x
-      )
-
-    result = f(x)
-    np.testing.assert_array_equal(x, result)
+    # TODO(b/395428868): Remove this check once we support subbyte types.
+    if jtu.test_device_matches(["tpu"]):
+      if "2" in dtype:
+        self.skipTest("TODO(dsuo): TPU callbacks send SIGABRT for int2/uint2.")
+      np.testing.assert_array_equal(jax.jit(f)(), np.arange(8, dtype=dtype))
+    else:
+      with self.assertRaisesRegex(
+          Exception, "Unsupported primitive type"
+      ):
+        _ = jax.jit(f)()
 
 
 class PureCallbackTest(jtu.JaxTestCase):
@@ -1136,6 +1108,20 @@ class PureCallbackTest(jtu.JaxTestCase):
       result += fun(jnp.ones((500, 500), jnp.complex64))[1]
     jax.block_until_ready(result)  # doesn't deadlock
 
+  def test_non_default_stride(self):
+    x = jnp.arange(24, dtype=jnp.float32).reshape(2, 3, 4)
+    def callback(x):
+      return np.asfortranarray(x)
+
+    @jax.jit
+    def f(x):
+      return jax.pure_callback(
+          callback, jax.ShapeDtypeStruct(x.shape, x.dtype), x
+      )
+
+    result = f(x)
+    np.testing.assert_array_equal(x, result)
+
 
 class IOCallbackTest(jtu.JaxTestCase):