diff --git a/.github/workflows/upstream-nightly.yml b/.github/workflows/upstream-nightly.yml
index 04df278019a5..e0dc14468871 100644
--- a/.github/workflows/upstream-nightly.yml
+++ b/.github/workflows/upstream-nightly.yml
@@ -85,7 +85,7 @@ jobs:
           && steps.status.outcome == 'failure'
           && github.event_name == 'schedule'
           && github.repository == 'jax-ml/jax'
-        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
+        uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b  # v4.5.0
         with:
           name: output-${{ matrix.python-version }}-log.jsonl
           path: output-${{ matrix.python-version }}-log.jsonl
diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
index 3904bf1b8f10..a9f498a816f4 100644
--- a/.github/workflows/wheel_win_x64.yml
+++ b/.github/workflows/wheel_win_x64.yml
@@ -45,7 +45,7 @@ jobs:
             --bazel_options=--config=win_clang `
             --verbose
 
-      - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
+      - uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b  # v4.5.0
         with:
           name: wheels-${{ matrix.os }}-${{ matrix.pyver }}
           path: ${{ github.workspace }}\dist\*.whl
diff --git a/.github/workflows/windows_ci.yml b/.github/workflows/windows_ci.yml
index 4c404ef4cb75..e94282ee133d 100644
--- a/.github/workflows/windows_ci.yml
+++ b/.github/workflows/windows_ci.yml
@@ -54,7 +54,7 @@ jobs:
             --bazel_options=--config=win_clang `
             --verbose
 
-      - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
+      - uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b  # v4.5.0
         with:
           name: wheels
           path: ${{ github.workspace }}\jax\dist\*.whl
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e86dece51013..bb9404268a3d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,10 @@ Best viewed [here](https://jax.readthedocs.io/en/latest/changelog.html).
 For the changes specific to the experimental Pallas APIs,
 see {ref}`pallas-changelog`.
 
+JAX follows Effort-based versioning; for a discussion of this and JAX's API
+compatibility policy, refer to {ref}`api-compatibility`. For the Python and
+NumPy version support policy, refer to {ref}`version-support-policy`.
+
 <!--
 Remember to align the itemized text with the first line of an item within a list.
 
@@ -12,30 +16,28 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 
 ## Unreleased
 
-* Deprecations
-  * From {mod}`jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
-    are now deprecated, having been replaced by symbols of the same name
-    in {mod}`jax.core`.
-
-* Deletions
-  * `jax_enable_memories` flag has been deleted and the behavior of that flag
-    is on by default.
-
 * Changes:
   * The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum
     supported version until June 2025.
 
-* Deprecations
-  * From {mod}`jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
-    are now deprecated, having been replaced by symbols of the same name
-    in {mod}`jax.core`.
-
 * New Features
   * {func}`jax.numpy.fft.fftn`, {func}`jax.numpy.fft.rfftn`,
     {func}`jax.numpy.fft.ifftn`, and {func}`jax.numpy.fft.irfftn` now support
     transforms in more than 3 dimensions, which was previously the limit. See
     {jax-issue}`#25606` for more details.
 
+* Deprecations
+  * From {mod}`jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
+    are now deprecated, having been replaced by symbols of the same name
+    in {mod}`jax.core`.
+
+* Deletions
+  * `jax_enable_memories` flag has been deleted and the behavior of that flag
+    is on by default.
+  * From `jax.lib.xla_client`, the previously-deprecated `Device` and
+    `XlaRuntimeError` symbols have been removed; instead use `jax.Device`
+    and `jax.errors.JaxRuntimeError` respectively.
+
 ## jax 0.4.38 (Dec 17, 2024)
 
 * Changes:
diff --git a/docs/profiling.md b/docs/profiling.md
index 91f4d61b21b6..ac992b3a05da 100644
--- a/docs/profiling.md
+++ b/docs/profiling.md
@@ -93,6 +93,12 @@ plugins" error described {ref}`below <multiple_installs>`. See
 <https://www.tensorflow.org/guide/profiler> for more information on installing
 TensorBoard.
 
+Nightly version of TensorBoard profiler requires nightly tensorflow and
+tensorboard
+```shell
+pip install tf-nightly tb-nightly tbp-nightly
+```
+
 ### Programmatic capture
 
 You can instrument your code to capture a profiler trace via the
diff --git a/docs/working-with-pytrees.md b/docs/working-with-pytrees.md
index 537a4df3e5a6..ffa47eba07c0 100644
--- a/docs/working-with-pytrees.md
+++ b/docs/working-with-pytrees.md
@@ -490,7 +490,7 @@ This section covers some of the most common patterns with JAX pytrees.
 
 ### Transposing pytrees with `jax.tree.map` and `jax.tree.transpose`
 
-To transpose a pytree (turn a list of trees into a tree of lists), JAX has two functions: {func} `jax.tree.map` (more basic) and {func}`jax.tree.transpose` (more flexible, complex and verbose).
+To transpose a pytree (turn a list of trees into a tree of lists), JAX has two functions: {func}`jax.tree.map` (more basic) and {func}`jax.tree.transpose` (more flexible, complex and verbose).
 
 **Option 1:** Use {func}`jax.tree.map`. Here's an example:
 
diff --git a/jax/_src/abstract_arrays.py b/jax/_src/abstract_arrays.py
index 2502b705b8fa..8ddc33fd8983 100644
--- a/jax/_src/abstract_arrays.py
+++ b/jax/_src/abstract_arrays.py
@@ -49,7 +49,6 @@ def masked_array_error(*args, **kwargs):
                    "Use arr.filled() to convert the value to a standard numpy array.")
 
 core.pytype_aval_mappings[np.ma.MaskedArray] = masked_array_error
-core.shaped_abstractify_handlers[np.ma.MaskedArray] = masked_array_error
 
 
 def _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:
@@ -58,7 +57,6 @@ def _make_shaped_array_for_numpy_array(x: np.ndarray) -> ShapedArray:
   return ShapedArray(x.shape, dtypes.canonicalize_dtype(dtype))
 
 core.pytype_aval_mappings[np.ndarray] = _make_shaped_array_for_numpy_array
-core.shaped_abstractify_handlers[np.ndarray] = _make_shaped_array_for_numpy_array
 
 
 def _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:
@@ -68,7 +66,6 @@ def _make_shaped_array_for_numpy_scalar(x: np.generic) -> ShapedArray:
 
 for t in numpy_scalar_types:
   core.pytype_aval_mappings[t] = _make_shaped_array_for_numpy_scalar
-  core.shaped_abstractify_handlers[t] = _make_shaped_array_for_numpy_scalar
 
 core.literalable_types.update(array_types)
 
@@ -81,6 +78,5 @@ def _make_abstract_python_scalar(typ, val):
 
 for t in dtypes.python_scalar_dtypes:
   core.pytype_aval_mappings[t] = partial(_make_abstract_python_scalar, t)
-  core.shaped_abstractify_handlers[t] = partial(_make_abstract_python_scalar, t)
 
 core.literalable_types.update(dtypes.python_scalar_dtypes.keys())
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 38ba4fd2d381..4bf964a72239 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2564,7 +2564,6 @@ def _sds_aval_mapping(x):
       x.shape, dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True),
       weak_type=x.weak_type)
 core.pytype_aval_mappings[ShapeDtypeStruct] = _sds_aval_mapping
-core.shaped_abstractify_handlers[ShapeDtypeStruct] = _sds_aval_mapping
 
 
 @api_boundary
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 1ce8e7786bb2..2ee8b01c77d4 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -1035,7 +1035,6 @@ def _get_aval_array(self):
   else:
     return self.aval
 
-core.shaped_abstractify_handlers[ArrayImpl] = _get_aval_array
 core.pytype_aval_mappings[ArrayImpl] = _get_aval_array
 
 # TODO(jakevdp) replace this with true inheritance at the C++ level.
diff --git a/jax/_src/compiler.py b/jax/_src/compiler.py
index 8ca7265bf2ae..5c42a3b44de3 100644
--- a/jax/_src/compiler.py
+++ b/jax/_src/compiler.py
@@ -192,6 +192,14 @@ def get_compile_options(
   build_options.exec_time_optimization_effort = config.exec_time_optimization_effort.value
   build_options.memory_fitting_effort = config.memory_fitting_effort.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:
+    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 = [
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 5f351bd46883..5d5173f3922a 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -656,6 +656,13 @@ def check_bool_conversion(arr: Array):
                      " is ambiguous. Use a.any() or a.all()")
 
 
+pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {}
+
+def _str_abstractify(x):
+  raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type")
+pytype_aval_mappings[str] = _str_abstractify
+
+
 def _aval_property(name):
   return property(lambda self: getattr(self.aval, name))
 
@@ -918,6 +925,8 @@ def unsafe_buffer_pointer(self):
 aval_property = namedtuple("aval_property", ["fget"])
 aval_method = namedtuple("aval_method", ["fun"])
 
+pytype_aval_mappings[Tracer] = lambda x: x.aval
+
 def check_eval_args(args):
   for arg in args:
     if isinstance(arg, Tracer):
@@ -1400,45 +1409,51 @@ def check_valid_jaxtype(x):
       f"Value {x!r} of type {type(x)} is not a valid JAX type")
 
 
-def _shaped_abstractify_slow(x):
-  try:
-    return x if isinstance(x, AbstractValue) else get_aval(x)
-  except TypeError:
-    pass
-
-  weak_type = getattr(x, 'weak_type', False)
-  if hasattr(x, 'dtype'):
-    dtype = dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True)
-  else:
-    raise TypeError(
-        f"Cannot interpret value of type {type(x)} as an abstract array; it "
-        "does not have a dtype attribute")
-  return ShapedArray(np.shape(x), dtype, weak_type=weak_type)
+# We have three flavors of abstractification APIs here which each used to have
+# their own separate implementation. Now they're effectively the same, with the
+# following differences:
+#
+# - abstractify returns avals for non-traced array-like objects.
+# - get_aval is like abstractify, but also accepts tracers.
+# - shaped_abstractify is like get_aval, but also accepts duck-typed arrays.
+#
+# TODO(jakevdp): can these be unified further?
 
-# TODO(jakevdp): deduplicate this with abstractify
 def shaped_abstractify(x):
-  # This was originally api_util.shaped_abstractify; temporarily moved
-  # here in order to facilitate combining it with abstractify.
-  handler = shaped_abstractify_handlers.get(type(x), None)
-  return handler(x) if handler is not None else _shaped_abstractify_slow(x)
+  typ = type(x)
+  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path
+    return aval_fn(x)
+  for t in typ.__mro__[1:]:
+    if (aval_fn := pytype_aval_mappings.get(t)):
+      return aval_fn(x)
+  if isinstance(x, AbstractValue):
+    return x
+  if hasattr(x, '__jax_array__'):
+    return shaped_abstractify(x.__jax_array__())
+  if hasattr(x, 'dtype'):
+    return ShapedArray(np.shape(x), x.dtype, weak_type=getattr(x, 'weak_type', False))
+  raise TypeError(
+      f"Cannot interpret value of type {typ} as an abstract array; it "
+      "does not have a dtype attribute")
 
 
 def abstractify(x):
-  for typ in type(x).__mro__:
-    aval_fn = pytype_aval_mappings.get(typ)
-    if aval_fn: return aval_fn(x)
-  if hasattr(x, '__jax_array__'):
-    return abstractify(x.__jax_array__())
-  raise TypeError(f"Argument '{x}' of type '{type(x)}' is not a valid JAX type")
+  if isinstance(x, Tracer):
+    raise TypeError(f"Argument '{x}' of type '{type(x)}' is not a valid JAX type")
+  return get_aval(x)
 
 
 def get_aval(x):
-  if isinstance(x, Tracer):
-    return x.aval
-  else:
-    return abstractify(x)
+  typ = type(x)
+  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path
+    return aval_fn(x)
+  for t in typ.__mro__[1:]:
+    if (aval_fn := pytype_aval_mappings.get(t)):
+      return aval_fn(x)
+  if hasattr(x, '__jax_array__'):
+    return get_aval(x.__jax_array__())
+  raise TypeError(f"Argument '{x}' of type '{typ}' is not a valid JAX type")
 
-get_type = get_aval
 
 def is_concrete(x):
   return to_concrete_value(x) is not None
@@ -1831,13 +1846,6 @@ def to_tangent_aval(self):
     return DShapedArray(self.shape, primal_dtype_to_tangent_dtype(self.dtype),
                         self.weak_type)
 
-pytype_aval_mappings: dict[type, Callable[[Any], AbstractValue]] = {}
-shaped_abstractify_handlers: dict[Any, Callable[[Any], AbstractValue]] = {}
-
-def _str_abstractify(x):
-  raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type")
-pytype_aval_mappings[str] = _str_abstractify
-shaped_abstractify_handlers[str] = _str_abstractify
 
 class DArray:
   _aval: DShapedArray
@@ -1894,7 +1902,6 @@ def _darray_aval(x):
   return DShapedArray(x._aval.shape, x._aval.dtype, x._aval.weak_type)
 
 pytype_aval_mappings[DArray] = _darray_aval
-shaped_abstractify_handlers[DArray] = _darray_aval
 
 
 @dataclass(frozen=True)
@@ -1924,11 +1931,10 @@ def __init__(self, aval, buf):
   aval = property(lambda self: self._aval)
   shape = property(lambda self: self._aval.shape)
   dtype = property(lambda self: self._aval.dtype)
-  def __getitem__(self, idx): return get_aval(self)._getitem(self, idx)
-  def __setitem__(self, idx, x): return get_aval(self)._setitem(self, idx, x)
+  def __getitem__(self, idx): return self._aval._getitem(self, idx)
+  def __setitem__(self, idx, x): return self._aval._setitem(self, idx, x)
   def __repr__(self) -> str: return 'Mutable' + repr(self[...])
 pytype_aval_mappings[MutableArray] = lambda x: x._aval
-shaped_abstractify_handlers[MutableArray] = lambda x: x._aval
 
 def mutable_array(init_val):
   return mutable_array_p.bind(init_val)
@@ -1984,7 +1990,6 @@ def __init__(self, buf):
   def block_until_ready(self):
     self._buf.block_until_ready()
 pytype_aval_mappings[Token] = lambda _: abstract_token
-shaped_abstractify_handlers[Token] = lambda _: abstract_token
 
 
 # TODO(dougalm): Deprecate these. They're just here for backwards compat.
diff --git a/jax/_src/cudnn/fused_attention_stablehlo.py b/jax/_src/cudnn/fused_attention_stablehlo.py
index a5a605002849..0963d762a1d1 100644
--- a/jax/_src/cudnn/fused_attention_stablehlo.py
+++ b/jax/_src/cudnn/fused_attention_stablehlo.py
@@ -348,11 +348,20 @@ def check_is_flash_attention(
             )
     else:
         # Regular attention conditions
-        if not ((H <= 128 and H % 8 == 0) and
-                (not is_training or not has_bias or T % 2 == 0 and S % 2 == 0)):
-            raise NotImplementedError(
-                f"Unsupported sequence length Q {T}, KV {S} and head dim {H}."
-            )
+        # Check the head dim.
+        is_on_hopper = check_compute_capability("9.0")
+        H_max = 256 if cudnn_version >= 90500 and is_on_hopper else 128
+        if not (H <= H_max and H % 8 == 0):
+          raise NotImplementedError(
+              f"The head dim must be <= {H_max} and a mutiple of 8, "
+              f"but got {H}."
+          )
+
+        # Check patterns with bias, seqlen should be divisible by 2
+        if (is_training and has_bias and (T % 2 != 0 or S % 2 != 0)):
+          raise NotImplementedError(
+              f"Unsupported sequence length Q {T}, KV {S}."
+          )
 
 def check_cudnn_version():
   # check if cuDNN is installed
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 8bde9c162036..769f9e44db61 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -300,8 +300,9 @@ def check_unused_args(self, used_args, args, kwargs):
 
 formatter = _DebugPrintFormatChecker()
 
-def _format_print_callback(fmt: str, *args, **kwargs):
-  sys.stdout.write(fmt.format(*args, **kwargs) + "\n")
+def _format_print_callback(fmt: str, np_printoptions, *args, **kwargs):
+  with np.printoptions(**np_printoptions):
+    sys.stdout.write(fmt.format(*args, **kwargs) + "\n")
 
 def debug_print(fmt: str, *args, ordered: bool = False, **kwargs) -> None:
   """Prints values and works in staged out JAX functions.
@@ -338,8 +339,8 @@ def debug_print(fmt: str, *args, **kwargs):
   # Check that we provide the correct arguments to be formatted.
   formatter.format(fmt, *args, **kwargs)
 
-  debug_callback(functools.partial(_format_print_callback, fmt), *args,
-                 **kwargs, ordered=ordered)
+  debug_callback(functools.partial(_format_print_callback, fmt, np.get_printoptions()),
+                 *args, **kwargs, ordered=ordered)
 
 
 # Sharding visualization
diff --git a/jax/_src/earray.py b/jax/_src/earray.py
index 25c2bc2bf7ec..98a0a863981e 100644
--- a/jax/_src/earray.py
+++ b/jax/_src/earray.py
@@ -115,7 +115,6 @@ def _earray_shard_arg_handler(xs, shardings, layouts, copy_semantics):
   return pxla.shard_args(phys_shardings, layouts, copy_semantics, arrs)
 pxla.shard_arg_handlers[EArray] = _earray_shard_arg_handler
 
-core.shaped_abstractify_handlers[EArray] = lambda self: self.aval
 core.pytype_aval_mappings[EArray] = lambda x: x.aval
 xla.canonicalize_dtype_handlers[EArray] = lambda x: x
 tree_util.dispatch_registry.register_node(
diff --git a/jax/_src/export/shape_poly.py b/jax/_src/export/shape_poly.py
index b82890cab682..5462723c8335 100644
--- a/jax/_src/export/shape_poly.py
+++ b/jax/_src/export/shape_poly.py
@@ -1205,7 +1205,6 @@ def _geq_decision(e1: DimSize, e2: DimSize, cmp_str: Callable[[], str]) -> bool:
       f"Symbolic dimension comparison {cmp_str()} is inconclusive.{describe_scope}")
 
 core.pytype_aval_mappings[_DimExpr] = _DimExpr._get_aval
-core.shaped_abstractify_handlers[_DimExpr] = _DimExpr._get_aval
 dtypes._weak_types.append(_DimExpr)
 
 def _convertible_to_int(p: DimSize) -> bool:
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 154b5e972682..ac0ae3a13967 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -1569,10 +1569,7 @@ def get_referent(self):
     val = frame.constvar_to_val.get(frame.tracer_to_var.get(id(self)))
     return self if val is None else get_referent(val)
 
-
-def _dynamic_jaxpr_tracer_shaped_abstractify(x):
-  return x.aval
-core.shaped_abstractify_handlers[DynamicJaxprTracer] = _dynamic_jaxpr_tracer_shaped_abstractify
+core.pytype_aval_mappings[DynamicJaxprTracer] = lambda x: x.aval
 
 def make_jaxpr_effects(constvars, invars, outvars, eqns) -> effects.Effects:
   sentinel = object()
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
index 301ebb181056..7566e6bf32fc 100644
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -710,7 +710,6 @@ def one_hot(x: Any, num_classes: int, *,
       'jax-nn-one-hot-float-input',
       f"jax.nn.one_hot input should be integer-typed; got dtype={x_arr.dtype}",
       stacklevel=1)
-    x_arr = x_arr.astype('int32')
   return _one_hot(x_arr, num_classes, dtype=dtype, axis=axis)
 
 
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 83ede1e48c3f..259c47948a9d 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -192,7 +192,6 @@ def __instancecheck__(self, instance: Any) -> bool:
 def _abstractify_scalar_meta(x):
   raise TypeError(f"JAX scalar type {x} cannot be interpreted as a JAX array.")
 core.pytype_aval_mappings[_ScalarMeta] = _abstractify_scalar_meta
-core.shaped_abstractify_handlers[_ScalarMeta] = _abstractify_scalar_meta
 
 def _make_scalar_type(np_scalar_type: type) -> _ScalarMeta:
   meta = _ScalarMeta(np_scalar_type.__name__, (object,),
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 02ca6f6ebab2..ff4e4e07e0e6 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -924,7 +924,7 @@ def pinv(a: ArrayLike, rtol: ArrayLike | None = None,
     - :func:`jax.numpy.linalg.inv`: multiplicative inverse of a square matrix.
 
   Notes:
-    :func:`jax.numpy.linalg.prng` differs from :func:`numpy.linalg.prng` in the
+    :func:`jax.numpy.linalg.pinv` differs from :func:`numpy.linalg.pinv` in the
     default value of `rcond``: in NumPy, the default  is `1e-15`. In JAX, the
     default is ``10. * max(num_rows, num_cols) * jnp.finfo(dtype).eps``.
 
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index de8688e491ba..10cfa5f7d22b 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -2765,6 +2765,12 @@ def log2(x: ArrayLike, /) -> Array:
     Array([-2., -1.,  0.,  1.,  2.,  3.], dtype=float32)
   """
   x, = promote_args_inexact("log2", x)
+  if dtypes.issubdtype(x.dtype, np.complexfloating):
+    r = lax.log(x)
+    re = lax.real(r)
+    im = lax.imag(r)
+    ln2 = lax.log(_constant_like(re, 2))
+    return lax.complex(lax.div(re, ln2), lax.div(im, ln2))
   return lax.div(lax.log(x), lax.log(_constant_like(x, 2)))
 
 
@@ -2789,6 +2795,12 @@ def log10(x: ArrayLike, /) -> Array:
     [-2. -1.  0.  1.  2.  3.]
   """
   x, = promote_args_inexact("log10", x)
+  if dtypes.issubdtype(x.dtype, np.complexfloating):
+    r = lax.log(x)
+    re = lax.real(r)
+    im = lax.imag(r)
+    ln10 = lax.log(_constant_like(re, 10))
+    return lax.complex(lax.div(re, ln10), lax.div(im, ln10))
   return lax.div(lax.log(x), lax.log(_constant_like(x, 10)))
 
 
diff --git a/jax/_src/pallas/mosaic/pallas_call_registration.py b/jax/_src/pallas/mosaic/pallas_call_registration.py
index ec9500c67cd7..8a7b64db6454 100644
--- a/jax/_src/pallas/mosaic/pallas_call_registration.py
+++ b/jax/_src/pallas/mosaic/pallas_call_registration.py
@@ -84,6 +84,8 @@ def _get_memory_space_from_aval(
       return None
     case tpu_core.TPUMemorySpace.VMEM:
       return tpu_custom_call.MemorySpace.VMEM
+    case tpu_core.TPUMemorySpace.SMEM:
+      return tpu_custom_call.MemorySpace.SMEM
     case tpu_core.TPUMemorySpace.SEMAPHORE:
       return tpu_custom_call.MemorySpace.SEMAPHORE_MEM
   return None
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index d492e4e1d86d..dc2d63f0b6ec 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -797,25 +797,30 @@ def _(step, carry):
     # Each range is 2 events, each event is 4 bytes.
     prof_spec = mgpu_profiler.ProfilerSpec(prof_space * 2 * 4)
     prof_ctx = ProfilerContext(params["profile_dir"], prof_spec)
-  module, out_structs_gmem, _ = mgpu_core._lower_as_gpu_kernel(
-      body,
-      grid=parallel_grid,
-      cluster=(),
-      block=block,
-      in_shapes=in_structs_gmem,
-      out_shape=out_structs_gmem,
-      smem_scratch_shape=(
-          (*in_structs_smem, *out_structs_smem),
-          *extra_smem_scratch,
-          (
-              mgpu.Barrier(arrival_count=1, num_barriers=max_concurrent_steps),
-              rs.barriers,
-              extra_barriers,
+  module, out_structs_gmem, _, launch_ctx, scratch_arr = (
+      mgpu_core._lower_as_gpu_kernel(
+          body,
+          grid=parallel_grid,
+          cluster=(),
+          block=block,
+          in_shapes=in_structs_gmem,
+          out_shape=out_structs_gmem,
+          smem_scratch_shape=(
+              (*in_structs_smem, *out_structs_smem),
+              *extra_smem_scratch,
+              (
+                  mgpu.Barrier(
+                      arrival_count=1, num_barriers=max_concurrent_steps
+                  ),
+                  rs.barriers,
+                  extra_barriers,
+              ),
           ),
-      ),
-      module_name=name_and_src_info.name,
-      prof_spec=prof_spec,
+          module_name=name_and_src_info.name,
+          prof_spec=prof_spec,
+      )
   )
+  mgpu_core._initialize_scratch(launch_ctx, scratch_arr)
 
   return LoweringResult(
       module, parallel_grid, block, out_structs_gmem, prof_ctx
@@ -1782,29 +1787,21 @@ def _ensure_fa(x: object, dtype: jnp.dtype) -> mgpu.FragmentedArray:
   if isinstance(x, mgpu.FragmentedArray):
     assert x.mlir_dtype == mgpu_utils.dtype_to_ir_type(dtype)
     return x
-  elif isinstance(x, (np.number, np.ndarray, int, float)):
-    return mgpu.FragmentedArray.splat(
-        _ir_constant(x, mgpu_utils.dtype_to_ir_type(dtype)),
-        (),
-        is_signed=mgpu_utils.is_signed(dtype),
-    )
-  elif isinstance(x, ir.Value):
-    if isinstance(x.type, (ir.IntegerType, ir.FloatType, ir.IndexType)):
-      assert x.type == mgpu_utils.dtype_to_ir_type(dtype)
-      return mgpu.FragmentedArray.splat(x, (), is_signed=mgpu_utils.is_signed(dtype))
-  raise NotImplementedError(f"Unsupported type: {type(x)}")
+  return mgpu.FragmentedArray.splat(
+      _ensure_ir_value(x, dtype), (), is_signed=mgpu_utils.is_signed(dtype)
+  )
 
 
 def _ensure_ir_value(x: object, dtype: jnp.dtype) -> ir.Value:
   if isinstance(x, ir.Value):
     assert x.type == mgpu_utils.dtype_to_ir_type(dtype)
     return x
-  elif isinstance(x, (np.number, np.ndarray, int, float)):
-    return _ir_constant(x, mgpu_utils.dtype_to_ir_type(dtype))
   elif isinstance(x, mgpu.FragmentedArray):
+    assert x.mlir_dtype == mgpu_utils.dtype_to_ir_type(dtype)
     if isinstance(x.layout, mgpu.WGSplatFragLayout):
       return x.registers.item()
-  raise NotImplementedError(f"Unsupported type: {type(x)}")
+    raise NotImplementedError(f"Unsupported layout: {x.layout}")
+  return _ir_constant(x, mgpu_utils.dtype_to_ir_type(dtype))
 
 
 def _ir_constant(v: object, t: ir.Type) -> ir.Value:
diff --git a/jax/_src/pallas/mosaic_gpu/primitives.py b/jax/_src/pallas/mosaic_gpu/primitives.py
index 85b7364ce2cc..9c053e2d0730 100644
--- a/jax/_src/pallas/mosaic_gpu/primitives.py
+++ b/jax/_src/pallas/mosaic_gpu/primitives.py
@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 import enum
 import math
 from typing import Any, Literal
@@ -25,7 +26,6 @@
 from jax._src import state
 from jax._src import tree_util
 from jax._src import util
-from jax._src.interpreters import mlir
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import arith as arith_dialect
 from jax._src.lib.mlir.dialects import llvm as llvm_dialect
@@ -33,23 +33,42 @@
 from jax._src.pallas import core as pallas_core
 from jax._src.pallas.mosaic_gpu import core as gpu_core
 from jax._src.pallas.mosaic_gpu import lowering
+from jax._src.pallas.mosaic_gpu.core import state_types
 from jax._src.state import discharge
 from jax._src.state import indexing
 from jax._src.state import primitives as state_primitives
-import jax.experimental.mosaic.gpu as mgpu
+from jax.experimental.mosaic import gpu as mgpu
+from jax.experimental.mosaic.gpu import utils as mgpu_utils
 import jax.numpy as jnp
 
 
 WARPGROUP_SIZE = 128
 
 
+_Ref = pallas_core.AbstractMemoryRef | state_types.TransformedRef
+
+
+def _check_ref(
+    aval: object, name: str, memory_space: gpu_core.GPUMemorySpace
+) -> None:
+  if not isinstance(aval, state_types.AbstractRef):
+    raise TypeError(f"{name} must be a reference, got {aval}")
+  aval_memory_space = getattr(aval, "memory_space", None) or gpu_core.GMEM
+  if aval_memory_space is not memory_space:
+    raise ValueError(
+        f"{name} must be a {memory_space.name.upper()} reference, got {aval}"
+    )
+
+
 copy_smem_to_gmem_p = jax_core.Primitive("copy_smem_to_gmem")
 copy_smem_to_gmem_p.multiple_results = True
 
 
 @copy_smem_to_gmem_p.def_effectful_abstract_eval
-def _copy_smem_to_gmem_abstract_eval(*avals, **params):
-  del avals, params  # Unused.
+def _copy_smem_to_gmem_abstract_eval(src, dst, *args, **params):
+  _check_ref(src, "src", gpu_core.SMEM)
+  _check_ref(dst, "dst", gpu_core.GMEM)
+  del args, params  # Unused.
   return (), {state.ReadEffect(0), state.WriteEffect(1)}
 
 
@@ -114,9 +133,7 @@ def _extract_smem_copy_params(transforms):
 
 
 def copy_smem_to_gmem(
-    src: pallas_core.AbstractMemoryRef,
-    dst: pallas_core.AbstractMemoryRef,
-    predicate: jax.Array | None = None,
+    src: _Ref, dst: _Ref, predicate: jax.Array | None = None
 ) -> None:
   """Asynchronously copies a SMEM reference to a GMEM reference.
 
@@ -130,10 +147,6 @@ def copy_smem_to_gmem(
     :func:`jax.experimental.mosaic.gpu.wait_smem_to_gmem`
     :func:`jax.experimental.mosaic.gpu.commit_smem`
   """
-  if src.memory_space is not gpu_core.SMEM:
-    raise TypeError(f"src must be a SMEM reference, got {src.memory_space}")
-  if getattr(dst, "memory_space", gpu_core.GMEM) is not gpu_core.GMEM:
-    raise ValueError(f"dst must be a GMEM reference, got {dst.memory_space}")
   src, src_transforms = state_primitives.get_ref_and_transforms(
       src, None, "copy_smem_to_gmem", force_trailing_indexer=False,
   )
@@ -164,8 +177,11 @@ def copy_smem_to_gmem(
 
 
 @copy_gmem_to_smem_p.def_effectful_abstract_eval
-def _copy_gmem_to_smem_abstract_eval(*avals, **params):
-  del avals, params  # Unused.
+def _copy_gmem_to_smem_abstract_eval(src, dst, barrier, *args, **params):
+  del args, params  # Unused.
+  _check_ref(src, "src", gpu_core.GMEM)
+  _check_ref(dst, "dst", gpu_core.SMEM)
+  _check_ref(barrier, "barrier", gpu_core.SMEM)
   return (), {state.ReadEffect(0), state.WriteEffect(1)}
 
 
@@ -217,21 +233,13 @@ def _copy_gmem_to_smem_lowering(
   return ()
 
 
-def copy_gmem_to_smem(
-    src: pallas_core.AbstractMemoryRef,
-    dst: pallas_core.AbstractMemoryRef,
-    barrier: pallas_core.AbstractMemoryRef,
-) -> None:
+def copy_gmem_to_smem(src: _Ref, dst: _Ref, barrier: _Ref) -> None:
   """Asynchronously copies a GMEM reference to a SMEM reference.
 
   See also:
     :func:`jax.experimental.mosaic.gpu.barrier_arrive`
     :func:`jax.experimental.mosaic.gpu.barrier_wait`
   """
-  if getattr(src, "memory_space", gpu_core.GMEM) is not gpu_core.GMEM:
-    raise TypeError(f"src must be a GMEM reference, got {src.memory_space}")
-  if dst.memory_space is not gpu_core.SMEM:
-    raise ValueError(f"dst must be a SMEM reference, got {dst.memory_space}")
   src, src_transforms = state_primitives.get_ref_and_transforms(
       src, None, "copy_gmem_to_smem", force_trailing_indexer=False,
   )
@@ -291,8 +299,9 @@ def _extract_barrier_indexer(transforms) -> indexing.NDIndexer | None:
 
 
 @barrier_arrive_p.def_effectful_abstract_eval
-def _barrier_arrive_abstract_eval(*avals, **params):
-  del avals, params  # Unused.
+def _barrier_arrive_abstract_eval(barrier, *args, **params):
+  del args, params  # Unused.
+  _check_ref(barrier, "barrier", gpu_core.SMEM)
   return (), {gpu_core._memory_effect}
 
 
@@ -328,8 +337,9 @@ def barrier_arrive(barrier: pallas_core.AbstractMemoryRef) -> None:
 
 
 @barrier_wait_p.def_effectful_abstract_eval
-def _barrier_wait_abstract_eval(*avals, **params):
-  del avals, params  # Unused.
+def _barrier_wait_abstract_eval(barrier, *args, **params):
+  _check_ref(barrier, "barrier", gpu_core.SMEM)
+  del args, params  # Unused.
   return (), {gpu_core._memory_effect}
 
 
@@ -703,38 +713,40 @@ def _broadcasted_iota_abstract_eval(dtype, shape, dimension, layout):
   del layout, dimension
   return jax_core.ShapedArray(shape, dtype)
 
-@lowering.register_lowering_rule(broadcasted_iota_p)
-def _broadcasted_iota_lowering(ctx: lowering.LoweringRuleContext, dtype, shape, dimension, layout):
-  del ctx
-  # Unsigned integers (as opposed to signless) cause MLIR verification
-  # errors so we only use signless like Mosaic GPU does.
-  #
-  # TODO(cperivol): use mgpu.utils.dtype_to_ir_type() instead.
-  mlir_dtype = (
-      ir.IntegerType.get_signless(dtype.itemsize * 8)
-      if jnp.issubdtype(dtype, jnp.integer)
-      else mlir.dtype_to_ir_type(dtype)
-  )
-  undef = llvm_dialect.mlir_undef(mlir_dtype)
-  is_signed = (
-      jnp.issubdtype(dtype, jnp.signedinteger)
-      if jnp.issubdtype(dtype, jnp.integer)
-      else None
-  )
 
-  i32 = ir.IntegerType.get_signless(32)
-  def _cast(x):
-    if ir.FloatType.isinstance(mlir_dtype):
-      x = arith_dialect.index_cast(i32, x)
-      return arith_dialect.uitofp(mlir_dtype, x)
-    else:
-      return arith_dialect.index_cast(mlir_dtype, x)
+@lowering.register_lowering_rule(broadcasted_iota_p)
+def _broadcasted_iota_lowering(
+    ctx: lowering.LoweringRuleContext, dtype, shape, dimension, layout
+):
+  del ctx  # Unused.
+  mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
+  if ir.FloatType.isinstance(mlir_dtype):
+    i32 = ir.IntegerType.get_signless(32)
+    cast = lambda x: arith_dialect.uitofp(
+        mlir_dtype, arith_dialect.index_cast(i32, x)
+    )
+  else:
+    cast = lambda x: arith_dialect.index_cast(mlir_dtype, x)
+  is_signed = mgpu_utils.is_signed(dtype)
   return mgpu.FragmentedArray.splat(
-      undef, shape, layout.value, is_signed=is_signed
+      llvm_dialect.mlir_undef(mlir_dtype),
+      shape,
+      layout.value,
+      is_signed=is_signed,
   ).foreach(
-      lambda _, idx: _cast(idx[dimension]), create_array=True, is_signed=is_signed
+      lambda _, idx: cast(idx[dimension]),
+      create_array=True,
+      is_signed=is_signed,
   )
 
 
-def broadcasted_iota(dtype, shape, dimension, *, layout: Layout | None = None):
-  return broadcasted_iota_p.bind(dtype=jnp.dtype(dtype), shape=shape, dimension=dimension, layout=layout)
+def broadcasted_iota(
+    dtype: jax.typing.DTypeLike,
+    shape: Sequence[int],
+    dimension: int,
+    *,
+    layout: Layout | None = None,
+) -> jax.Array:
+  return broadcasted_iota_p.bind(
+      dtype=jnp.dtype(dtype), shape=shape, dimension=dimension, layout=layout
+  )
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
index d29bad5d5304..4f43b54bb478 100644
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -461,8 +461,6 @@ def __hash__(self) -> int:
 
 
 core.pytype_aval_mappings[PRNGKeyArray] = lambda x: x.aval
-core.shaped_abstractify_handlers[PRNGKeyArray] = op.attrgetter('aval')
-
 xla.canonicalize_dtype_handlers[PRNGKeyArray] = lambda x: x
 
 
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index 10963330ff92..2720ca196012 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -249,8 +249,8 @@ def as_text(self) -> str:
       else:
         raise
 
-    # TODO(skyewm): this should return a single dict (I think returning a list
-    # was to support MPMD executables, which never fully landed)
+    # TODO(b/384741132): this should return a single dict (I think returning a list
+    # was to support MPMD executables, which never fully landed).
   def cost_analysis(self) -> list[dict[str, float]]:
     xla_ext_exe = self.xla_extension_executable()
 
@@ -266,9 +266,19 @@ def cost_analysis(self) -> list[dict[str, float]]:
     # Try client method if executable cost_analysis method is unimplemented
     if hasattr(xla_ext_exe, "client"):
       try:
+        # TODO(b/384741132): We expect that the executable has only one
+        # HloModule. We should be able to remove this check once we update the
+        # Executable class to have only a single HloModule (see bug).
+        hlo_modules = xla_ext_exe.hlo_modules()
+        assert len(hlo_modules) == 1, (
+            f"Exectuable should have only one HloModule ({len(hlo_modules)})"
+            " were found)."
+        )
+
         return [
-            xla_extension.hlo_module_cost_analysis(xla_ext_exe.client, m)
-            for m in xla_ext_exe.hlo_modules()
+            xla_extension.hlo_module_cost_analysis(
+                xla_ext_exe.client, hlo_modules[0]
+            )
         ]
       except xla_extension.XlaRuntimeError as e:
         msg, *_ = e.args
diff --git a/jax/_src/tpu_custom_call.py b/jax/_src/tpu_custom_call.py
index 2d77acba02da..bb92afebe8e9 100644
--- a/jax/_src/tpu_custom_call.py
+++ b/jax/_src/tpu_custom_call.py
@@ -83,6 +83,7 @@ class MemorySpace(enum.Enum):
   HBM = enum.auto()
   VMEM = enum.auto()
   SEMAPHORE_MEM = enum.auto()
+  SMEM = enum.auto()
 
   @property
   def color(self) -> int:
@@ -92,6 +93,8 @@ def color(self) -> int:
       return 1
     elif self == MemorySpace.SEMAPHORE_MEM:
       return 2
+    elif self == MemorySpace.SMEM:
+      return 4
     else:
       raise ValueError("invalid memory space: " + str(self))
 
diff --git a/jax/core.py b/jax/core.py
index 54bbdac51c87..ef1551b2f1ba 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -128,7 +128,7 @@
                              _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_type),
+    "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.",
@@ -212,7 +212,7 @@
   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_type
+  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
diff --git a/jax/experimental/rnn.py b/jax/experimental/rnn.py
index 4aa863708189..f13b4fe029bf 100644
--- a/jax/experimental/rnn.py
+++ b/jax/experimental/rnn.py
@@ -175,6 +175,31 @@ def init_lstm_weight(rng: PRNGKeyArray, input_size: int, hidden_size: int,
   return jax.random.uniform(
       rng, shape=(param_count,), dtype=jnp.float32, minval=-k, maxval=k)
 
+def swap_lstm_gates(weights, input_size, hidden_size, num_layers, bidirectional):
+  """Swaps the weights for the input and output gates for an LSTM model."""
+  weights = jnp.asarray(weights)  # Ensure weights are JAX arrays
+  flat_shapes = _get_params_shapes_in_lstm(input_size, hidden_size, num_layers, bidirectional)
+  num_directions = 2 if bidirectional else 1
+
+  w_offsets = 0
+  for l in range(num_layers):
+    for direction in range(num_directions):
+      # Iterate through all weight and bias gate names to swap gates in both weights and biases
+      for gate_name in ["W_ih", "W_hh", "b_ih", "b_hh"]:
+        shape = flat_shapes.pop(0)  # Get the current shape and remove it from the list
+        num_elems = math.prod(shape)
+        matrix = weights[w_offsets:w_offsets + num_elems].reshape(shape)
+
+        # Swap between the input and output gates (third and fourth gates)
+        gates = jnp.split(matrix, 4, axis=0)
+        swapped_matrix = jnp.concatenate([gates[0], gates[1], gates[3], gates[2]], axis=0)
+
+        # Update the weights with swapped matrix
+        weights = weights.at[w_offsets:w_offsets + num_elems].set(swapped_matrix.flatten())
+        w_offsets += num_elems
+
+  return weights
+
 
 def unpack_lstm_weights(
     weights: Array, input_size: int, hidden_size: int, num_layers: int,
@@ -437,7 +462,8 @@ def _gpu_lowering_strip_tf32(fn, *args, cudnn_allow_tf32, **kw):
 rnn_fwd_p.def_impl(partial(xla.apply_primitive, rnn_fwd_p))
 rnn_fwd_p.def_abstract_eval(rnn_abstract_eval)
 if gpu_rnn:
-  mlir.register_lowering(rnn_fwd_p, gpu_rnn.cudnn_rnn_lowering, platform='cuda')
+  mlir.register_lowering(rnn_fwd_p, gpu_rnn.cudnn_rnn_fwd_lowering, platform='cuda')
+  mlir.register_lowering(rnn_fwd_p, gpu_rnn.miopen_rnn_fwd_lowering, platform='rocm')
 
 
 def lstm_bwd(input_size: int, hidden_size: int, num_layers: int, dropout: float,
@@ -481,5 +507,7 @@ def rnn_bwd_abstract_eval(dy_aval, dhn_aval, dcn_aval, x_aval, h0_aval, c0_aval,
 if gpu_rnn:
   mlir.register_lowering(
       rnn_bwd_p, gpu_rnn.cudnn_rnn_bwd_lowering, platform='cuda')
+  mlir.register_lowering(
+      rnn_bwd_p, gpu_rnn.miopen_rnn_bwd_lowering, platform='rocm')
 
 lstm.defvjp(lstm_fwd, lstm_bwd)
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index 572440f486eb..f98415dbac82 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -724,7 +724,6 @@ def _xla_shard(ctx: mlir.LoweringRuleContext, mesh, auto, names,
                aval_in, aval_out, x):
   if prod([size for n, size in mesh.shape.items() if n not in auto]) == 1:
     return x
-  manual_proto = pxla.manual_proto(aval_in, frozenset(mesh.axis_names) - auto, mesh)
   axes = {name: i for i, ns in names.items() for name in ns}
   ns = _make_scoped_manual_sharding(ctx, mesh, axes)
   if dtypes.issubdtype(aval_in.dtype, dtypes.extended):
@@ -734,6 +733,7 @@ def _xla_shard(ctx: mlir.LoweringRuleContext, mesh, auto, names,
   unspecified = set(range(aval_in.ndim)) if auto else set()
   sx = mlir.wrap_with_sharding_op(ctx, x, aval_in, shard_proto,
                                   unspecified_dims=unspecified)
+  manual_proto = pxla.manual_proto(aval_in, frozenset(mesh.axis_names) - auto, mesh)
   return mlir.wrap_with_full_to_shard_op(ctx, sx, aval_out, manual_proto, unspecified)
 
 def _xla_unshard(ctx: mlir.LoweringRuleContext, mesh, auto, names,
@@ -746,6 +746,8 @@ def _xla_unshard(ctx: mlir.LoweringRuleContext, mesh, auto, names,
     ns = sharding_impls.physical_sharding(aval_out, ns)
     aval_out = core.physical_aval(aval_out)
   unspecified = set(range(aval_out.ndim)) if auto else set()
+  if dtypes.issubdtype(aval_in.dtype, dtypes.extended):
+    aval_in = core.physical_aval(aval_in)
   manual_proto = pxla.manual_proto(aval_in, frozenset(mesh.axis_names) - auto, mesh)
   sx = mlir.wrap_with_sharding_op(ctx, x, aval_in, manual_proto, unspecified_dims=unspecified)
   shard_proto = ns._to_xla_hlo_sharding(aval_out.ndim).to_proto()
diff --git a/jax/lib/xla_client.py b/jax/lib/xla_client.py
index f6c4d16e6db5..86e7307c804b 100644
--- a/jax/lib/xla_client.py
+++ b/jax/lib/xla_client.py
@@ -36,17 +36,17 @@
         "jax.lib.xla_client.bfloat16 was removed in JAX v0.4.38; use ml_dtypes.bfloat16.",
         None,
     ),
-    # Added Sep 26 2024
+    # Finalized 2024-12-23; remove after 2024-03-23
     "Device": (
         "jax.lib.xla_client.Device is deprecated; use jax.Device instead.",
-        _xc.Device,
+        None,
     ),
     "XlaRuntimeError": (
         (
             "jax.lib.xla_client.XlaRuntimeError is deprecated; use"
             " jax.errors.JaxRuntimeError."
         ),
-        _xc.XlaRuntimeError,
+        None,
     ),
     # Added Oct 10 2024
     "FftType": (
diff --git a/jaxlib/gpu/rnn_kernels.cc b/jaxlib/gpu/rnn_kernels.cc
index 27fb8f9c4a06..80e00c27a2a7 100644
--- a/jaxlib/gpu/rnn_kernels.cc
+++ b/jaxlib/gpu/rnn_kernels.cc
@@ -30,7 +30,7 @@ namespace jax {
 namespace JAX_GPU_NAMESPACE {
 
 std::string ErrorString(gpudnnStatus_t status) {
-  return cudnnGetErrorString(status);
+  return gpudnnGetErrorString(status);
 }
 
 template <typename T>
@@ -80,63 +80,88 @@ DoRnnComputeWorkspaceReserveSpaceSizes(int input_size, int hidden_size,
   JAX_RETURN_IF_ERROR(h.status());
   auto& handle = *h;
 
-  cudnnRNNDescriptor_t rnn_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateRNNDescriptor(&rnn_desc)));
+  gpudnnRNNDescriptor_t rnn_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateRNNDescriptor(&rnn_desc)));
 
-  cudnnDropoutDescriptor_t dropout_desc;
+  gpudnnDropoutDescriptor_t dropout_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateDropoutDescriptor(&dropout_desc)));
+      JAX_AS_STATUS(gpudnnCreateDropoutDescriptor(&dropout_desc)));
   size_t state_size;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDropoutGetStatesSize(handle.get(), &state_size)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetDropoutDescriptor(
+      JAX_AS_STATUS(gpudnnDropoutGetStatesSize(handle.get(), &state_size)));
+
+#ifdef JAX_GPU_HIP
+  void* dropout_states_dev = nullptr;
+  // Allocate minimal memory for dropout states (can be very small since it's not used)
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(hipMalloc(&dropout_states_dev, state_size)));
+  if (!dropout_states_dev) {
+    return absl::InternalError("Failed to allocate minimal GPU memory for dropout states.");
+  }
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
+      dropout_desc, handle.get(), dropout, dropout_states_dev, state_size, 123, false, false,
+      MIOPEN_RNG_PSEUDO_XORWOW)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
       dropout_desc, handle.get(), dropout, nullptr, state_size, 123)));
+#endif // JAX_GPU_HIP
 
   // TODO(zhangqiaorjc): Handle other kinds of RNN.
-  cudnnRNNMode_t cell_mode = CUDNN_LSTM;
-  cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
+  gpudnnRNNMode_t cell_mode = GPUDNN_LSTM;
+  gpudnnRNNBiasMode_t bias_mode = GPUDNN_RNN_DOUBLE_BIAS;
   int num_directions = 1;
-  cudnnDirectionMode_t dir_mode = CUDNN_UNIDIRECTIONAL;
+  gpudnnDirectionMode_t dir_mode = GPUDNN_UNIDIRECTIONAL;
   if (bidirectional) {
-    dir_mode = CUDNN_BIDIRECTIONAL;
+    dir_mode = GPUDNN_BIDIRECTIONAL;
     num_directions = 2;
   }
-  cudnnRNNInputMode_t input_mode = CUDNN_LINEAR_INPUT;
-  cudnnDataType_t data_type = CUDNN_DATA_FLOAT;
-  cudnnDataType_t math_prec = CUDNN_DATA_FLOAT;
-  cudnnMathType_t math_type = cudnn_allow_tf32? CUDNN_DEFAULT_MATH: CUDNN_FMA_MATH;
+  gpudnnRNNInputMode_t input_mode = GPUDNN_LINEAR_INPUT;
+  gpudnnDataType_t data_type = GPUDNN_DATA_FLOAT;
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, hidden_size, num_layers, dropout_desc, input_mode, dir_mode,
+      cell_mode, bias_mode, GPUDNN_RNN_ALGO_STANDARD, data_type)));
+#else // JAX_GPU_CUDA
+  gpudnnDataType_t math_prec = GPUDNN_DATA_FLOAT;
+  gpudnnMathType_t math_type = cudnn_allow_tf32? GPUDNN_DEFAULT_MATH: GPUDNN_FMA_MATH;
   int32_t proj_size = hidden_size;
-  uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDescriptor_v8(
-      rnn_desc, CUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
+  uint32_t aux_flags = GPUDNN_RNN_PADDED_IO_ENABLED;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, GPUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
       input_mode, data_type, math_prec, math_type, input_size, hidden_size,
       proj_size, num_layers, dropout_desc, aux_flags)));
+#endif // JAX_GPU_HIP
 
-  cudnnForwardMode_t fwdMode = CUDNN_FWD_MODE_TRAINING;
-  cudnnRNNDataLayout_t layout = CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
+  gpudnnForwardMode_t fwdMode = GPUDNN_FWD_MODE_TRAINING;
+  gpudnnRNNDataLayout_t layout = GPUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
   float padding = 0.0f;
 
   std::vector<int32_t> seq_length_vector(batch_size, max_seq_length);
   int32_t* seq_length_array = &seq_length_vector[0];
 
-  cudnnRNNDataDescriptor_t input_data_desc;
+  gpudnnRNNDataDescriptor_t input_data_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateRNNDataDescriptor(&input_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDataDescriptor(
+      JAX_AS_STATUS(gpudnnCreateRNNDataDescriptor(&input_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDataDescriptor(
       input_data_desc, data_type, layout, max_seq_length, batch_size,
       input_size, seq_length_array, &padding)));
 
   size_t workSpaceSize;
   size_t reserveSpaceSize;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnGetRNNTempSpaceSizes(
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnGetRNNTempSpaceSizes(
+      handle.get(), rnn_desc, input_data_desc, fwdMode, &workSpaceSize,
+      &reserveSpaceSize)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnGetRNNTempSpaceSizes(
       handle.get(), rnn_desc, fwdMode, input_data_desc, &workSpaceSize,
       &reserveSpaceSize)));
-
+#endif // JAX_GPU_HIP
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyDropoutDescriptor(dropout_desc)));
+      JAX_AS_STATUS(gpudnnDestroyDropoutDescriptor(dropout_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyRNNDataDescriptor(input_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyRNNDescriptor(rnn_desc)));
+      JAX_AS_STATUS(gpudnnDestroyRNNDataDescriptor(input_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyRNNDescriptor(rnn_desc)));
 
   // Round up to nearest multiples of 4 so we can return them as f32 arrays.
   workSpaceSize += (workSpaceSize % 4);
@@ -162,41 +187,61 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers,
   JAX_RETURN_IF_ERROR(h.status());
   auto& handle = *h;
 
-  cudnnRNNDescriptor_t rnn_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateRNNDescriptor(&rnn_desc)));
+  gpudnnRNNDescriptor_t rnn_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateRNNDescriptor(&rnn_desc)));
 
-  cudnnDropoutDescriptor_t dropout_desc;
+  gpudnnDropoutDescriptor_t dropout_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateDropoutDescriptor(&dropout_desc)));
+      JAX_AS_STATUS(gpudnnCreateDropoutDescriptor(&dropout_desc)));
   size_t state_size;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDropoutGetStatesSize(handle.get(), &state_size)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetDropoutDescriptor(
+      JAX_AS_STATUS(gpudnnDropoutGetStatesSize(handle.get(), &state_size)));
+
+  void* dropout_states_dev = nullptr;
+  // Allocate minimal memory for dropout states (can be very small since it's not used).
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(hipMalloc(&dropout_states_dev, state_size)));
+  if (!dropout_states_dev) {
+    return absl::InternalError("Failed to allocate minimal GPU memory for dropout states.");
+  }
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
+      dropout_desc, handle.get(), d.dropout, dropout_states_dev, state_size, 123, false, false,
+      MIOPEN_RNG_PSEUDO_XORWOW)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
       dropout_desc, handle.get(), d.dropout, nullptr, state_size, 123)));
+#endif // JAX_GPU_HIP
 
   // TODO(zhangqiaorjc): Handle other kinds of RNN.
-  cudnnRNNMode_t cell_mode = CUDNN_LSTM;
-  cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
+  gpudnnRNNMode_t cell_mode = GPUDNN_LSTM;
+  gpudnnRNNBiasMode_t bias_mode = GPUDNN_RNN_DOUBLE_BIAS;
   int num_directions = 1;
-  cudnnDirectionMode_t dir_mode = CUDNN_UNIDIRECTIONAL;
+  gpudnnDirectionMode_t dir_mode = GPUDNN_UNIDIRECTIONAL;
   if (d.bidirectional) {
-    dir_mode = CUDNN_BIDIRECTIONAL;
+    dir_mode = GPUDNN_BIDIRECTIONAL;
     num_directions = 2;
   }
-  cudnnRNNInputMode_t input_mode = CUDNN_LINEAR_INPUT;
-  cudnnDataType_t data_type = CUDNN_DATA_FLOAT;
-  cudnnDataType_t math_prec = CUDNN_DATA_FLOAT;
-  cudnnMathType_t math_type = d.cudnn_allow_tf32? CUDNN_DEFAULT_MATH: CUDNN_FMA_MATH;
+  gpudnnRNNInputMode_t input_mode = GPUDNN_LINEAR_INPUT;
+  gpudnnDataType_t data_type = GPUDNN_DATA_FLOAT;
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, d.hidden_size, d.num_layers, dropout_desc, input_mode, dir_mode,
+      cell_mode, bias_mode, GPUDNN_RNN_ALGO_STANDARD, data_type)));
+#else // JAX_GPU_CUDA
+  gpudnnDataType_t math_prec = GPUDNN_DATA_FLOAT;
+  gpudnnMathType_t math_type = d.cudnn_allow_tf32? GPUDNN_DEFAULT_MATH: GPUDNN_FMA_MATH;
   int32_t proj_size = d.hidden_size;
-  uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;
-
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDescriptor_v8(
-      rnn_desc, CUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
+  uint32_t aux_flags = GPUDNN_RNN_PADDED_IO_ENABLED;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, GPUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
       input_mode, data_type, math_prec, math_type, d.input_size, d.hidden_size,
       proj_size, d.num_layers, dropout_desc, aux_flags)));
+#endif // JAX_GPU_HIP
 
-  cudnnForwardMode_t fwdMode = CUDNN_FWD_MODE_TRAINING;
-  cudnnRNNDataLayout_t layout = CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
+  gpudnnForwardMode_t fwdMode = GPUDNN_FWD_MODE_TRAINING;
+  gpudnnRNNDataLayout_t layout = GPUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
   float padding = 0.0f;
 
   // TODO(zhangqiaorjc): Avoid this cudaMemcpy if possible.
@@ -209,17 +254,17 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers,
                                    gpuMemcpyDeviceToHost, stream)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuStreamSynchronize(stream)));
 
-  cudnnRNNDataDescriptor_t input_data_desc;
+  gpudnnRNNDataDescriptor_t input_data_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateRNNDataDescriptor(&input_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDataDescriptor(
+      JAX_AS_STATUS(gpudnnCreateRNNDataDescriptor(&input_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDataDescriptor(
       input_data_desc, data_type, layout, d.max_seq_length, d.batch_size,
       d.input_size, seq_length_array, &padding)));
 
-  cudnnRNNDataDescriptor_t output_data_desc;
+  gpudnnRNNDataDescriptor_t output_data_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateRNNDataDescriptor(&output_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDataDescriptor(
+      JAX_AS_STATUS(gpudnnCreateRNNDataDescriptor(&output_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDataDescriptor(
       output_data_desc, data_type, layout, d.max_seq_length, d.batch_size,
       d.hidden_size * num_directions, seq_length_array, &padding)));
 
@@ -232,19 +277,31 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers,
   strides[0] = dims[1] * dims[2];
   strides[1] = dims[2];
   strides[2] = 1;
-  cudnnTensorDescriptor_t h_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateTensorDescriptor(&h_desc)));
+  gpudnnTensorDescriptor_t h_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateTensorDescriptor(&h_desc)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnSetTensorNdDescriptor(h_desc, data_type, 3, dims, strides)));
+      gpudnnSetTensorNdDescriptor(h_desc, data_type, 3, dims, strides)));
 
-  cudnnTensorDescriptor_t c_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateTensorDescriptor(&c_desc)));
+  gpudnnTensorDescriptor_t c_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateTensorDescriptor(&c_desc)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnSetTensorNdDescriptor(c_desc, data_type, 3, dims, strides)));
+      gpudnnSetTensorNdDescriptor(c_desc, data_type, 3, dims, strides)));
 
   size_t weight_space_size;
+#ifdef JAX_GPU_HIP
+miopenTensorDescriptor_t input_tensor_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(miopenCreateTensorDescriptor(&input_tensor_desc)));
+  int dimsA[2] = {d.batch_size, d.input_size};
+  int stridesA[2] = {dimsA[1], 1}; // Row-major order, similar to GPUDNN
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(miopenSetTensorDescriptor(
+    input_tensor_desc, data_type, 2, dimsA, stridesA)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
+      gpudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, input_tensor_desc,
+      &weight_space_size, data_type)));
+#else // JAX_GPU_CUDA
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, &weight_space_size)));
+      gpudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, &weight_space_size)));
+#endif // JAX_GPU_HIP
 
   auto input_buf = buffers[0];
   auto h_0_buf = buffers[1];
@@ -255,22 +312,30 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers,
   auto c_n_buf = buffers[7];
   auto workspace_buf = buffers[8];
   auto reserve_space_buf = buffers[9];
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnRNNForward(
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNForward(
+      handle.get(), rnn_desc, fwdMode, input_data_desc, input_buf,
+      h_desc, h_0_buf, h_n_buf, c_desc, c_0_buf, c_n_buf,
+      output_data_desc, output_buf, weights_buf, weight_space_size,
+      workspace_buf, d.workspace_size, reserve_space_buf, d.reserve_space_size)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNForward(
       handle.get(), rnn_desc, fwdMode, (const int32_t*)seq_lengths_buf,
       input_data_desc, input_buf, output_data_desc, output_buf, h_desc, h_0_buf,
       h_n_buf, c_desc, c_0_buf, c_n_buf, weight_space_size, weights_buf,
-      d.workspace_size, workspace_buf, d.reserve_space_size,
-      reserve_space_buf)));
+      workspace_buf, d.workspace_size, reserve_space_buf, d.reserve_space_size)));
+#endif // JAX_GPU_HIP
 
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyTensorDescriptor(h_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyTensorDescriptor(c_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(h_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(c_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyDropoutDescriptor(dropout_desc)));
+      JAX_AS_STATUS(gpudnnDestroyRNNDataDescriptor(input_data_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyRNNDataDescriptor(input_data_desc)));
+      JAX_AS_STATUS(gpudnnDestroyRNNDataDescriptor(output_data_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyRNNDataDescriptor(output_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyRNNDescriptor(rnn_desc)));
+      JAX_AS_STATUS(gpudnnDestroyDropoutDescriptor(dropout_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyRNNDescriptor(rnn_desc)));
 
   return absl::OkStatus();
 }
@@ -284,40 +349,60 @@ static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers,
   JAX_RETURN_IF_ERROR(h.status());
   auto& handle = *h;
 
-  cudnnRNNDescriptor_t rnn_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateRNNDescriptor(&rnn_desc)));
+  gpudnnRNNDescriptor_t rnn_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateRNNDescriptor(&rnn_desc)));
 
-  cudnnDropoutDescriptor_t dropout_desc;
+  gpudnnDropoutDescriptor_t dropout_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateDropoutDescriptor(&dropout_desc)));
+      JAX_AS_STATUS(gpudnnCreateDropoutDescriptor(&dropout_desc)));
   size_t state_size;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDropoutGetStatesSize(handle.get(), &state_size)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetDropoutDescriptor(
+      JAX_AS_STATUS(gpudnnDropoutGetStatesSize(handle.get(), &state_size)));
+
+  void* dropout_states_dev = nullptr;
+  // Allocate minimal memory for dropout states (can be very small since it's not used)
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(hipMalloc(&dropout_states_dev, state_size)));
+  if (!dropout_states_dev) {
+    return absl::InternalError("Failed to allocate minimal GPU memory for dropout states.");
+  }
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
+      dropout_desc, handle.get(), d.dropout, dropout_states_dev, state_size, 123, false, false,
+      MIOPEN_RNG_PSEUDO_XORWOW)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
       dropout_desc, handle.get(), d.dropout, nullptr, state_size, 123)));
+#endif // JAX_GPU_HIP
 
   // TODO(zhangqiaorjc): Handle other kinds of RNN.
-  cudnnRNNMode_t cell_mode = CUDNN_LSTM;
-  cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
+  gpudnnRNNMode_t cell_mode = GPUDNN_LSTM;
+  gpudnnRNNBiasMode_t bias_mode = GPUDNN_RNN_DOUBLE_BIAS;
   int num_directions = 1;
-  cudnnDirectionMode_t dir_mode = CUDNN_UNIDIRECTIONAL;
+  gpudnnDirectionMode_t dir_mode = GPUDNN_UNIDIRECTIONAL;
   if (d.bidirectional) {
-    dir_mode = CUDNN_BIDIRECTIONAL;
+    dir_mode = GPUDNN_BIDIRECTIONAL;
     num_directions = 2;
   }
-  cudnnRNNInputMode_t input_mode = CUDNN_LINEAR_INPUT;
-  cudnnDataType_t data_type = CUDNN_DATA_FLOAT;
-  cudnnDataType_t math_prec = CUDNN_DATA_FLOAT;
-  cudnnMathType_t math_type = d.cudnn_allow_tf32? CUDNN_DEFAULT_MATH: CUDNN_FMA_MATH;
+  gpudnnRNNInputMode_t input_mode = GPUDNN_LINEAR_INPUT;
+  gpudnnDataType_t data_type = GPUDNN_DATA_FLOAT;
+
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, d.hidden_size, d.num_layers, dropout_desc, input_mode, dir_mode,
+      cell_mode, bias_mode, GPUDNN_RNN_ALGO_STANDARD, data_type)));
+#else // JAX_GPU_CUDA
+  gpudnnDataType_t math_prec = GPUDNN_DATA_FLOAT;
+  gpudnnMathType_t math_type = d.cudnn_allow_tf32? GPUDNN_DEFAULT_MATH: GPUDNN_FMA_MATH;
   int32_t proj_size = d.hidden_size;
-  uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;
-
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDescriptor_v8(
-      rnn_desc, CUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
+  uint32_t aux_flags = GPUDNN_RNN_PADDED_IO_ENABLED;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDescriptor(
+      rnn_desc, GPUDNN_RNN_ALGO_STANDARD, cell_mode, bias_mode, dir_mode,
       input_mode, data_type, math_prec, math_type, d.input_size, d.hidden_size,
       proj_size, d.num_layers, dropout_desc, aux_flags)));
+#endif // JAX_GPU_HIP
 
-  cudnnRNNDataLayout_t layout = CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
+  gpudnnRNNDataLayout_t layout = GPUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
   float padding = 0.0f;
 
   auto seq_lengths_buf = buffers[10];
@@ -329,17 +414,17 @@ static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers,
                                    gpuMemcpyDeviceToHost, stream)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuStreamSynchronize(stream)));
 
-  cudnnRNNDataDescriptor_t input_data_desc;
+  gpudnnRNNDataDescriptor_t input_data_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateRNNDataDescriptor(&input_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDataDescriptor(
+      JAX_AS_STATUS(gpudnnCreateRNNDataDescriptor(&input_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDataDescriptor(
       input_data_desc, data_type, layout, d.max_seq_length, d.batch_size,
       d.input_size, seq_length_array, &padding)));
 
-  cudnnRNNDataDescriptor_t output_data_desc;
+  gpudnnRNNDataDescriptor_t output_data_desc;
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnCreateRNNDataDescriptor(&output_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnSetRNNDataDescriptor(
+      JAX_AS_STATUS(gpudnnCreateRNNDataDescriptor(&output_data_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetRNNDataDescriptor(
       output_data_desc, data_type, layout, d.max_seq_length, d.batch_size,
       d.hidden_size * num_directions, seq_length_array, &padding)));
 
@@ -352,19 +437,31 @@ static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers,
   strides[0] = dims[1] * dims[2];
   strides[1] = dims[2];
   strides[2] = 1;
-  cudnnTensorDescriptor_t h_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateTensorDescriptor(&h_desc)));
+  gpudnnTensorDescriptor_t h_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateTensorDescriptor(&h_desc)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnSetTensorNdDescriptor(h_desc, data_type, 3, dims, strides)));
+      gpudnnSetTensorNdDescriptor(h_desc, data_type, 3, dims, strides)));
 
-  cudnnTensorDescriptor_t c_desc;
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnCreateTensorDescriptor(&c_desc)));
+  gpudnnTensorDescriptor_t c_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnCreateTensorDescriptor(&c_desc)));
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnSetTensorNdDescriptor(c_desc, data_type, 3, dims, strides)));
+      gpudnnSetTensorNdDescriptor(c_desc, data_type, 3, dims, strides)));
 
   size_t weight_space_size;
+#ifdef JAX_GPU_HIP
+  miopenTensorDescriptor_t input_tensor_desc;
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(miopenCreateTensorDescriptor(&input_tensor_desc)));
+  int input_dims[2] = {d.batch_size, d.input_size};
+  int input_strides[2] = {input_dims[1], 1}; // row-major order
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(miopenSetTensorDescriptor(
+    input_tensor_desc, data_type, 2, input_dims, input_strides)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
+      gpudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, input_tensor_desc,
+      &weight_space_size, data_type)));
+#else // JAX_GPU_CUDA
   JAX_RETURN_IF_ERROR(JAX_AS_STATUS(
-      cudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, &weight_space_size)));
+      gpudnnGetRNNWeightSpaceSize(handle.get(), rnn_desc, &weight_space_size)));
+#endif // JAX_GPU_HIP
 
   auto dy_buf = buffers[0];
   auto dh_n_buf = buffers[1];
@@ -384,29 +481,43 @@ static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers,
   // auto dw_buf = buffers[14];
   auto workspace_buf = buffers[15];
 
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnRNNBackwardData_v8(
+#ifdef JAX_GPU_HIP
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNBackwardData(
+      handle.get(), rnn_desc, output_data_desc, y_buf, dy_buf,
+      h_desc, h_0_buf, dh_n_buf, dh_0_buf,
+      c_desc, c_0_buf, dc_n_buf, dc_0_buf,
+      input_data_desc, dx_buf, w_buf, weight_space_size,
+      workspace_buf, d.workspace_size, reserve_space_buf, d.reserve_space_size)));
+
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNBackwardWeights(
+      handle.get(), rnn_desc, input_data_desc, x_buf, h_desc, h_0_buf,
+      output_data_desc, y_buf, zeroed_dw_buf, weight_space_size,
+      workspace_buf, d.workspace_size, reserve_space_buf, d.reserve_space_size)));
+#else // JAX_GPU_CUDA
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNBackwardData(
       handle.get(), rnn_desc, (const int32_t*)seq_lengths_buf, output_data_desc,
       y_buf, dy_buf, input_data_desc, dx_buf, h_desc, h_0_buf, dh_n_buf,
       dh_0_buf, c_desc, c_0_buf, dc_n_buf, dc_0_buf, weight_space_size, w_buf,
       d.workspace_size, workspace_buf, d.reserve_space_size,
       reserve_space_buf)));
 
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnRNNBackwardWeights_v8(
-      handle.get(), rnn_desc, CUDNN_WGRAD_MODE_ADD,
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnRNNBackwardWeights(
+      handle.get(), rnn_desc, GPUDNN_WGRAD_MODE_ADD,
       (const int32_t*)seq_lengths_buf, input_data_desc, x_buf, h_desc, h_0_buf,
       output_data_desc, y_buf, weight_space_size, zeroed_dw_buf,
       d.workspace_size, workspace_buf, d.reserve_space_size,
       reserve_space_buf)));
+#endif // JAX_GPU_HIP
 
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyTensorDescriptor(h_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyTensorDescriptor(c_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(h_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(c_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyDropoutDescriptor(dropout_desc)));
+      JAX_AS_STATUS(gpudnnDestroyRNNDataDescriptor(input_data_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyRNNDataDescriptor(input_data_desc)));
+      JAX_AS_STATUS(gpudnnDestroyRNNDataDescriptor(output_data_desc)));
   JAX_RETURN_IF_ERROR(
-      JAX_AS_STATUS(cudnnDestroyRNNDataDescriptor(output_data_desc)));
-  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(cudnnDestroyRNNDescriptor(rnn_desc)));
+      JAX_AS_STATUS(gpudnnDestroyDropoutDescriptor(dropout_desc)));
+  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyRNNDescriptor(rnn_desc)));
 
   return absl::OkStatus();
 }
diff --git a/jaxlib/gpu/vendor.h b/jaxlib/gpu/vendor.h
index 648580f08a92..a74e1b79a75d 100644
--- a/jaxlib/gpu/vendor.h
+++ b/jaxlib/gpu/vendor.h
@@ -72,6 +72,18 @@ typedef CUevent gpuEvent_t;
 typedef CUfunction gpuFunction_t;
 typedef cudnnHandle_t gpudnnHandle_t;
 typedef cudnnStatus_t gpudnnStatus_t;
+typedef cudnnRNNDescriptor_t gpudnnRNNDescriptor_t;
+typedef cudnnDropoutDescriptor_t gpudnnDropoutDescriptor_t;
+typedef cudnnTensorDescriptor_t gpudnnTensorDescriptor_t;
+typedef cudnnRNNDataDescriptor_t gpudnnRNNDataDescriptor_t;
+typedef cudnnRNNDataLayout_t gpudnnRNNDataLayout_t;
+typedef cudnnMathType_t gpudnnMathType_t;
+typedef cudnnDataType_t gpudnnDataType_t;
+typedef cudnnRNNInputMode_t gpudnnRNNInputMode_t;
+typedef cudnnDirectionMode_t gpudnnDirectionMode_t;
+typedef cudnnRNNBiasMode_t gpudnnRNNBiasMode_t;
+typedef cudnnRNNMode_t gpudnnRNNMode_t;
+typedef cudnnForwardMode_t gpudnnForwardMode_t;
 typedef CUmodule gpuModule_t;
 typedef cusolverDnHandle_t gpusolverDnHandle_t;
 typedef cusolverStatus_t gpusolverStatus_t;
@@ -114,9 +126,41 @@ typedef cusparseDnVecDescr_t gpusparseDnVecDescr_t;
 #define GPUBLAS_STATUS_SUCCESS CUBLAS_STATUS_SUCCESS
 
 #define gpudnnCreate cudnnCreate
+#define gpudnnGetErrorString cudnnGetErrorString
+#define gpudnnCreateRNNDescriptor cudnnCreateRNNDescriptor
 #define gpudnnSetStream cudnnSetStream
+#define gpudnnDropoutGetStatesSize cudnnDropoutGetStatesSize
+#define gpudnnSetDropoutDescriptor cudnnSetDropoutDescriptor
+#define gpudnnDestroyRNNDescriptor cudnnDestroyRNNDescriptor
+#define gpudnnDestroyRNNDataDescriptor cudnnDestroyRNNDataDescriptor
+#define gpudnnDestroyTensorDescriptor cudnnDestroyTensorDescriptor
+#define gpudnnDestroyDropoutDescriptor cudnnDestroyDropoutDescriptor
+#define gpudnnRNNBackwardWeights cudnnRNNBackwardWeights_v8
+#define gpudnnRNNBackwardData cudnnRNNBackwardData_v8
+#define gpudnnGetRNNWeightSpaceSize cudnnGetRNNWeightSpaceSize
+#define gpudnnCreateTensorDescriptor cudnnCreateTensorDescriptor
+#define gpudnnSetTensorNdDescriptor cudnnSetTensorNdDescriptor
+#define gpudnnCreateRNNDataDescriptor cudnnCreateRNNDataDescriptor
+#define gpudnnSetRNNDataDescriptor cudnnSetRNNDataDescriptor
+#define gpudnnSetRNNDescriptor cudnnSetRNNDescriptor_v8
+#define gpudnnCreateDropoutDescriptor cudnnCreateDropoutDescriptor
+#define gpudnnGetRNNTempSpaceSizes cudnnGetRNNTempSpaceSizes
+#define gpudnnRNNForward cudnnRNNForward
 
 #define GPUDNN_STATUS_SUCCESS CUDNN_STATUS_SUCCESS
+#define GPUDNN_WGRAD_MODE_ADD CUDNN_WGRAD_MODE_ADD
+#define GPUDNN_RNN_ALGO_STANDARD CUDNN_RNN_ALGO_STANDARD
+#define GPUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED
+#define GPUDNN_RNN_PADDED_IO_ENABLED CUDNN_RNN_PADDED_IO_ENABLED
+#define GPUDNN_DEFAULT_MATH CUDNN_DEFAULT_MATH
+#define GPUDNN_FMA_MATH CUDNN_FMA_MATH
+#define GPUDNN_DATA_FLOAT CUDNN_DATA_FLOAT
+#define GPUDNN_LINEAR_INPUT CUDNN_LINEAR_INPUT
+#define GPUDNN_FWD_MODE_TRAINING CUDNN_FWD_MODE_TRAINING
+#define GPUDNN_UNIDIRECTIONAL CUDNN_UNIDIRECTIONAL
+#define GPUDNN_RNN_DOUBLE_BIAS CUDNN_RNN_DOUBLE_BIAS
+#define GPUDNN_LSTM CUDNN_LSTM
+#define GPUDNN_BIDIRECTIONAL CUDNN_BIDIRECTIONAL
 
 #define gpusolverDnCreate cusolverDnCreate
 #define gpusolverDnSetStream cusolverDnSetStream
@@ -364,6 +408,7 @@ constexpr uint32_t kNumThreadsPerWarp = 32;
 #include "rocm/include/hipblas/hipblas.h"
 #include "rocm/include/hipsolver/hipsolver.h"
 #include "rocm/include/hipsparse/hipsparse.h"
+#include "rocm/include/miopen/miopen.h"
 // IWYU pragma: end_exports
 
 #define JAX_GPU_NAMESPACE hip
@@ -372,6 +417,9 @@ constexpr uint32_t kNumThreadsPerWarp = 32;
 #define JAX_GPU_HAVE_SPARSE 1
 #define JAX_GPU_HAVE_64_BIT 0
 #define JAX_GPU_HAVE_FP8 0
+// TODO(Ruturaj4): Currently equivalent API does exist in
+// MIOpen lib. Remove when MIOpen support is complete.
+#define MIOPEN_STATUS_SUCCESS 0
 
 typedef hipFloatComplex gpuComplex;
 typedef hipDoubleComplex gpuDoubleComplex;
@@ -394,6 +442,19 @@ typedef hipStream_t gpuStream_t;
 typedef hipError_t gpuError_t;
 typedef hipEvent_t gpuEvent_t;
 typedef hipFunction_t gpuFunction_t;
+typedef miopenHandle_t gpudnnHandle_t;
+typedef miopenStatus_t gpudnnStatus_t;
+typedef miopenRNNDescriptor_t gpudnnRNNDescriptor_t;
+typedef miopenDropoutDescriptor_t gpudnnDropoutDescriptor_t;
+typedef miopenTensorDescriptor_t gpudnnTensorDescriptor_t;
+typedef miopenSeqTensorDescriptor_t gpudnnRNNDataDescriptor_t;
+typedef miopenRNNBaseLayout_t gpudnnRNNDataLayout_t;
+typedef miopenDataType_t gpudnnDataType_t;
+typedef miopenRNNInputMode_t gpudnnRNNInputMode_t;
+typedef miopenRNNDirectionMode_t gpudnnDirectionMode_t;
+typedef miopenRNNBiasMode_t gpudnnRNNBiasMode_t;
+typedef miopenRNNMode_t gpudnnRNNMode_t;
+typedef miopenRNNFWDMode_t gpudnnForwardMode_t;
 typedef hipModule_t gpuModule_t;
 typedef void gpuSyevjInfo;
 typedef hipsolverSyevjInfo_t gpuSyevjInfo_t;
@@ -432,6 +493,39 @@ typedef hipsparseDnVecDescr_t gpusparseDnVecDescr_t;
 
 #define GPUBLAS_STATUS_SUCCESS HIPBLAS_STATUS_SUCCESS
 
+#define gpudnnCreate miopenCreate
+#define gpudnnGetErrorString miopenGetErrorString
+#define gpudnnSetStream miopenSetStream
+#define gpudnnCreateRNNDescriptor miopenCreateRNNDescriptor
+#define gpudnnDropoutGetStatesSize miopenDropoutGetStatesSize
+#define gpudnnSetDropoutDescriptor miopenSetDropoutDescriptor
+#define gpudnnDestroyRNNDescriptor miopenDestroyRNNDescriptor
+#define gpudnnDestroyRNNDataDescriptor miopenDestroySeqTensorDescriptor
+#define gpudnnDestroyTensorDescriptor miopenDestroyTensorDescriptor
+#define gpudnnDestroyDropoutDescriptor miopenDestroyDropoutDescriptor
+#define gpudnnRNNBackwardWeights miopenRNNBackwardWeightsSeqTensor
+#define gpudnnCreateRNNDataDescriptor miopenCreateSeqTensorDescriptor
+#define gpudnnRNNBackwardData miopenRNNBackwardSeqData
+#define gpudnnCreateTensorDescriptor miopenCreateTensorDescriptor
+#define gpudnnSetTensorNdDescriptor miopenSetTensorDescriptor
+#define gpudnnSetRNNDataDescriptor miopenSetRNNDataSeqTensorDescriptor
+#define gpudnnSetRNNDescriptor miopenSetRNNDescriptor_V2
+#define gpudnnCreateDropoutDescriptor miopenCreateDropoutDescriptor
+#define gpudnnGetRNNTempSpaceSizes miopenGetRNNTempSpaceSizes
+#define gpudnnRNNForward miopenRNNForward
+#define gpudnnGetRNNWeightSpaceSize miopenGetRNNParamsSize
+
+#define GPUDNN_STATUS_SUCCESS MIOPEN_STATUS_SUCCESS
+#define GPUDNN_RNN_ALGO_STANDARD miopenRNNdefault
+#define GPUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED miopenRNNDataSeqMajorPadded
+#define GPUDNN_DATA_FLOAT miopenFloat
+#define GPUDNN_LINEAR_INPUT miopenRNNlinear
+#define GPUDNN_FWD_MODE_TRAINING miopenRNNTraining
+#define GPUDNN_UNIDIRECTIONAL miopenRNNunidirection
+#define GPUDNN_RNN_DOUBLE_BIAS miopenRNNwithBias
+#define GPUDNN_LSTM miopenLSTM
+#define GPUDNN_BIDIRECTIONAL miopenRNNbidirection
+
 #define gpusolverDnCreate hipsolverCreate
 #define gpusolverDnSetStream hipsolverSetStream
 #define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo
diff --git a/jaxlib/gpu_rnn.py b/jaxlib/gpu_rnn.py
index 0fc3dc350967..9bb3bd5b2756 100644
--- a/jaxlib/gpu_rnn.py
+++ b/jaxlib/gpu_rnn.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import partial
 import importlib
 
 import jaxlib.mlir.ir as ir
@@ -24,19 +25,33 @@
 
 for cuda_module_name in [".cuda", "jax_cuda12_plugin"]:
   try:
-    _rnn = importlib.import_module(f"{cuda_module_name}._rnn", package="jaxlib")
+    _cuda_rnn = importlib.import_module(f"{cuda_module_name}._rnn", package="jaxlib")
   except ImportError:
-    _rnn = None
+    _cuda_rnn = None
   else:
     break
 
-if _rnn:
-  for _name, _value in _rnn.registrations().items():
+if _cuda_rnn:
+  for _name, _value in _cuda_rnn.registrations().items():
     xla_client.register_custom_call_target(_name, _value, platform='CUDA')
-  compute_rnn_workspace_reserve_space_sizes = _rnn.compute_rnn_workspace_reserve_space_sizes
+  compute_rnn_workspace_reserve_space_sizes = _cuda_rnn.compute_rnn_workspace_reserve_space_sizes
 
 
-def cudnn_rnn_lowering(ctx, input, h_0, c_0, weights, seq_lengths, *,
+for rocm_module_name in [".rocm", "jax_rocm60_plugin"]:
+  try:
+    _hip_rnn = importlib.import_module(f"{rocm_module_name}._rnn", package="jaxlib")
+  except ImportError:
+    _hip_rnn = None
+  else:
+    break
+
+if _hip_rnn:
+  for _name, _value in _hip_rnn.registrations().items():
+    xla_client.register_custom_call_target(_name, _value, platform='ROCM')
+  compute_rnn_workspace_reserve_space_sizes = _hip_rnn.compute_rnn_workspace_reserve_space_sizes
+
+
+def _rnn_fwd_lowering(_rnn, platform, ctx, input, h_0, c_0, weights, seq_lengths, *,
                        input_size: int, hidden_size: int, num_layers: int,
                        dropout: bool, bidirectional: bool,
                        cudnn_allow_tf32: bool):
@@ -75,11 +90,10 @@ def cudnn_rnn_lowering(ctx, input, h_0, c_0, weights, seq_lengths, *,
                                      reserve_space_shape[0])
 
   i32_type = ir.IntegerType.get_signless(32)
-
   out = hlo.CustomCallOp(
       [output_type, h_0.type, c_0.type, workspace_type, reserve_space_type],
       [input, h_0, c_0, weights, seq_lengths],
-      call_target_name=ir.StringAttr.get('cudnn_rnn'),
+      call_target_name=ir.StringAttr.get(f"{platform}dnn_rnn"),
       has_side_effect=ir.BoolAttr.get(False),
       backend_config=ir.StringAttr.get(opaque),
       api_version=ir.IntegerAttr.get(i32_type, 2),
@@ -87,6 +101,10 @@ def cudnn_rnn_lowering(ctx, input, h_0, c_0, weights, seq_lengths, *,
   )
   return out.results[:-2] + out.results[-1:]  # drop workspace output
 
+cudnn_rnn_fwd_lowering = partial(_rnn_fwd_lowering, _cuda_rnn, "cu")
+
+miopen_rnn_fwd_lowering = partial(_rnn_fwd_lowering, _hip_rnn, "hip")
+
 
 def _hlo_zeros_f32(shape):
   return hlo.constant(
@@ -94,7 +112,7 @@ def _hlo_zeros_f32(shape):
           np.zeros(shape, dtype=np.float32), type=ir.F32Type.get()))
 
 
-def cudnn_rnn_bwd_lowering(ctx, dy, dhn, dcn, x, h0, c0, w, y,
+def _rnn_bwd_lowering(_rnn, platform, ctx, dy, dhn, dcn, x, h0, c0, w, y,
                            reserve_space, seq_lengths, *, input_size: int,
                            hidden_size: int, num_layers: int, dropout: bool,
                            bidirectional: bool, cudnn_allow_tf32: bool):
@@ -123,7 +141,7 @@ def cudnn_rnn_bwd_lowering(ctx, dy, dhn, dcn, x, h0, c0, w, y,
           dy, dhn, dcn, x, h0, c0, w, y, reserve_space, zeroed_dw,
           seq_lengths
       ],
-      call_target_name=ir.StringAttr.get('cudnn_rnn_bwd'),
+      call_target_name=ir.StringAttr.get(f"{platform}dnn_rnn_bwd"),
       has_side_effect=ir.BoolAttr.get(False),
       backend_config=ir.StringAttr.get(opaque),
       api_version=ir.IntegerAttr.get(i32_type, 2),
@@ -135,3 +153,7 @@ def cudnn_rnn_bwd_lowering(ctx, dy, dhn, dcn, x, h0, c0, w, y,
               operand_tuple_indices=[])
       ]))
   return out.results[:-1]  # drop workspace output
+
+cudnn_rnn_bwd_lowering = partial(_rnn_bwd_lowering, _cuda_rnn, "cu")
+
+miopen_rnn_bwd_lowering = partial(_rnn_bwd_lowering, _hip_rnn, "hip")
diff --git a/jaxlib/mlir/_mlir_libs/BUILD.bazel b/jaxlib/mlir/_mlir_libs/BUILD.bazel
index 511a2310b3b4..4486a4c4fca2 100644
--- a/jaxlib/mlir/_mlir_libs/BUILD.bazel
+++ b/jaxlib/mlir/_mlir_libs/BUILD.bazel
@@ -70,8 +70,8 @@ py_extension(
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPIGPUHeaders",
         "@llvm-project//mlir:CAPIIRHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -85,7 +85,8 @@ py_extension(
     deps = [
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPIGPUHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -100,8 +101,8 @@ py_extension(
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:CAPINVGPUHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -116,8 +117,8 @@ py_extension(
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:CAPILLVMHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -131,8 +132,8 @@ py_extension(
     deps = [
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPISparseTensorHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -146,7 +147,8 @@ py_extension(
     deps = [
         ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//mlir:CAPISparseTensorHeaders",
-        "@pybind11",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
+        "@nanobind",
     ],
 )
 
@@ -156,9 +158,10 @@ py_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi",
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi_headers",
         "@llvm-project//mlir:CAPIIRHeaders",
-        "@llvm-project//mlir:MLIRBindingsPythonHeadersAndDeps",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
         "@nanobind",
     ],
 )
@@ -378,6 +381,7 @@ cc_library(
     name = "jaxlib_mlir_capi_objects",
     deps = [
         "//jaxlib/mosaic:tpu_dialect_capi_objects",
+        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi_objects",
         "@llvm-project//mlir:CAPIArithObjects",
         "@llvm-project//mlir:CAPIGPUObjects",
         "@llvm-project//mlir:CAPIIRObjects",
diff --git a/jaxlib/mosaic/BUILD b/jaxlib/mosaic/BUILD
index 238bf42d9651..37f9a35596d6 100644
--- a/jaxlib/mosaic/BUILD
+++ b/jaxlib/mosaic/BUILD
@@ -28,6 +28,7 @@ package(
 py_library(
     name = "mosaic",
     deps = [
+        "//jaxlib/mosaic/python:gpu_dialect",
         "//jaxlib/mosaic/python:tpu_dialect",
     ],
 )
@@ -42,6 +43,7 @@ cc_library(
         "dialect/tpu/tpu_dialect.cc",
         "dialect/tpu/tpu_ops.cc",
         "dialect/tpu/util.cc",
+        "dialect/tpu/vreg_util.cc",
         ":extension_srcs",
     ] + glob([
         "dialect/tpu/transforms/*.cc",
@@ -50,6 +52,7 @@ cc_library(
         "dialect/tpu/layout.h",
         "dialect/tpu/tpu_dialect.h",
         "dialect/tpu/util.h",
+        "dialect/tpu/vreg_util.h",
     ] + glob([
         "dialect/tpu/transforms/*.h",
     ]),
@@ -231,6 +234,19 @@ cc_library(
     alwayslink = True,
 )
 
+cc_test(
+    name = "vreg_util_test",
+    srcs = ["dialect/tpu/vreg_util_test.cc"],
+    deps = [
+        ":tpu_dialect",
+        "//testing/base/public:gunit_main",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:VectorDialect",
+    ],
+)
+
 filegroup(
     name = "extension_srcs",
     srcs = [
diff --git a/jaxlib/mosaic/dialect/gpu/BUILD b/jaxlib/mosaic/dialect/gpu/BUILD
index 681ee708edd8..50ea58104459 100644
--- a/jaxlib/mosaic/dialect/gpu/BUILD
+++ b/jaxlib/mosaic/dialect/gpu/BUILD
@@ -215,3 +215,26 @@ cc_library(
         "@llvm-project//mlir:CAPIIR",
     ],
 )
+
+# Header-only target, used when using the C API from a separate shared library.
+cc_library(
+    name = "gpu_dialect_capi_headers",
+    hdrs = DIALECT_CAPI_HEADERS,
+    deps = [
+        ":mosaic_gpu_inc_gen",
+        "@llvm-project//mlir:CAPIIRHeaders",
+    ],
+)
+
+# Alwayslink target, used when exporting the C API from a shared library.
+cc_library(
+    name = "gpu_dialect_capi_objects",
+    srcs = DIALECT_CAPI_SOURCES,
+    hdrs = DIALECT_CAPI_HEADERS,
+    deps = [
+        ":mosaic_gpu",
+        ":mosaic_gpu_inc_gen",
+        "@llvm-project//mlir:CAPIIRObjects",
+    ],
+    alwayslink = True,
+)
diff --git a/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc b/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
index 34f6241661d5..e2e1b623b624 100644
--- a/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
+++ b/jaxlib/mosaic/dialect/gpu/mosaic_gpu_test.cc
@@ -19,8 +19,8 @@ limitations under the License.
 #include <optional>
 #include <string>
 
-#include "testing/base/public/gmock.h"
-#include "testing/base/public/gunit.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
index 613e02460141..3a8263573544 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -29,7 +29,6 @@
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -52,6 +51,7 @@
 #include "absl/status/status.h"
 #include "absl/types/span.h"
 #include "llvm/include/llvm/ADT/APInt.h"
+#include "llvm/include/llvm/Support/LogicalResult.h"
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/include/mlir/Dialect/Vector/IR/VectorOps.h"
@@ -64,6 +64,7 @@
 #include "jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout_extensions.h"
 #include "jaxlib/mosaic/dialect/tpu/transforms/infer_memref_layout.h"
 #include "jaxlib/mosaic/dialect/tpu/util.h"
+#include "jaxlib/mosaic/dialect/tpu/vreg_util.h"
 #include "xla/array.h"
 #include "xla/layout.h"
 #include "xla/util.h"
@@ -275,16 +276,6 @@ void updateSliceFromRange(xla::Array<T> &arr, Range data,
   CHECK(data_it == data.end());
 }
 
-FailureOr<TypedAttr> getZeroIntOrFloatAttr(Type ty) {
-  if (isa<FloatType>(ty)) {
-    return TypedAttr(FloatAttr::get(ty, 0));
-  }
-  if (isa<IntegerType>(ty)) {
-    return TypedAttr(IntegerAttr::get(ty, 0));
-  }
-  return emitError(UnknownLoc::get(ty.getContext()), "Not implemented: ") << ty;
-}
-
 FailureOr<int64_t> getIntConst(Value v, bool silent = false) {
   if (auto constant_op = v.getDefiningOp<arith::ConstantOp>()) {
     if (auto integer_attr = dyn_cast<IntegerAttr>(constant_op.getValue())) {
@@ -479,33 +470,6 @@ FailureOr<BlockArgument> appendConstant(RewriteContext &ctx, func::FuncOp func,
   return argument;
 }
 
-VectorType getNativeVregOrVmaskTypeImpl(
-    Type elem_ty, const int8_t bitwidth,
-    const std::array<int64_t, 2> target_shape) {
-  if (bitwidth == 32) {
-    return VectorType::get(target_shape, elem_ty);
-  }
-  return VectorType::get({target_shape[0], target_shape[1], 32 / bitwidth},
-                         elem_ty);
-}
-
-VectorType getNativeVregOrVmaskType(Type elem_ty, const int8_t layout_bitwidth,
-                                    const std::array<int64_t, 2> target_shape) {
-  int8_t bitwidth = elem_ty.getIntOrFloatBitWidth();
-  if (bitwidth == 1) {
-    bitwidth = layout_bitwidth;
-  } else {
-    CHECK_EQ(bitwidth, layout_bitwidth);
-  }
-  return getNativeVregOrVmaskTypeImpl(elem_ty, bitwidth, target_shape);
-}
-
-VectorType getNativeVregType(Type elem_ty,
-                             const std::array<int64_t, 2> target_shape) {
-  return getNativeVregOrVmaskTypeImpl(elem_ty, elem_ty.getIntOrFloatBitWidth(),
-                                      target_shape);
-}
-
 // Masks all values outside of bounds.
 //
 // Arguments:
@@ -518,7 +482,7 @@ VectorType getNativeVregType(Type elem_ty,
 // Returns:
 //   An MLIR value of the same type as the value argument, with all entries
 //   outside of bounds replaced by neutral.
-FailureOr<Value> maskOOB(RewriteContext &ctx, OpBuilder &builder,
+FailureOr<Value> maskOOB(RewriteContext &ctx, ImplicitLocOpBuilder &builder,
                          TypedValue<VectorType> value,
                          const VRegDataBounds &bounds,
                          const Attribute neutral) {
@@ -542,9 +506,7 @@ FailureOr<Value> maskOOB(RewriteContext &ctx, OpBuilder &builder,
         value.getLoc(),
         VectorType::get(native_vreg_ty.getShape(), builder.getI1Type()), mask);
   }
-  auto neutral_vec = builder.create<arith::ConstantOp>(
-      value.getLoc(), native_vreg_ty,
-      DenseElementsAttr::get(native_vreg_ty, neutral));
+  Value neutral_vec = getFullVector(builder, native_vreg_ty, neutral);
   return builder
       .create<arith::SelectOp>(value.getLoc(), mask, value, neutral_vec)
       .getResult();
@@ -1863,126 +1825,28 @@ LogicalResult tpu_matmul_rule(RewriteContext &ctx, Operation &op,
   TPU_ASSERT_EQ_OP(padded_lhs_rows, lhs_vregs.dim(0) * layout_lhs.tiling()[0]);
   TPU_ASSERT_EQ_OP(padded_rhs_rows, rhs_vregs.dim(0) * layout_rhs.tiling()[0]);
 
-  const VectorType i32_vreg_ty =
-      getNativeVregType(builder.getI32Type(), ctx.target_shape);
-  auto getX32VmaskByPaddingEnd = [&](int64_t dim, int64_t padding) {
-    CHECK(dim == 0 || dim == 1);
-    CHECK(padding >= 0 && padding <= ctx.target_shape[dim]);
-    return cast<TypedValue<VectorType>>(
-        builder
-            .create<arith::CmpIOp>(
-                arith::CmpIPredicate::slt,
-                builder.create<tpu::IotaOp>(i32_vreg_ty,
-                                            builder.getI32IntegerAttr(dim)),
-                builder.create<arith::ConstantOp>(DenseElementsAttr::get(
-                    i32_vreg_ty, builder.getI32IntegerAttr(
-                                     ctx.target_shape[dim] - padding))))
-            .getResult());
-  };
-
-  // We can also extend this helper function with padding_top and padding_left
-  // based on the offsets in vregs.
-  const Value i32_zeros_vreg = builder.create<arith::ConstantOp>(
-      op.getLoc(),
-      DenseElementsAttr::get(i32_vreg_ty, builder.getI32IntegerAttr(0)));
-  const Value i32_max_vreg = builder.create<arith::ConstantOp>(
-      op.getLoc(), DenseElementsAttr::get(
-                       i32_vreg_ty, builder.getI32IntegerAttr(0xffffffff)));
-  auto maskVregs = [&](xla::Array<Value> &vregs, int64_t padding_bottom,
-                       int64_t padding_right) {
-    auto vreg_ty = cast<VectorType>(vregs.begin()->getType());
-    int packing = vreg_ty.getRank() > 2 ? vreg_ty.getShape()[2] : 1;
-    // Mask out the bottom.
-    if (padding_bottom > 0) {
-      // We have limited the row size of LHS and RHS need to be a multiple of
-      // native tiling at the beginning of this rule. Therefore, it is safe to
-      // bitcast to x32 vreg for masking.
-      int sub_padding = padding_bottom % packing;
-      int x32_padding_bottom = padding_bottom / packing;
-      auto mask_bottom = getX32VmaskByPaddingEnd(0, x32_padding_bottom);
-      // Create an int32 vreg which contains subelement masking and then
-      // logical_and with target vreg to mask out the unaligned paddings.
-      // Eg. if padding_bottom = 5, packing = 2, and assume the vreg shape is
-      // [8, 128], then the mask will be:
-      //
-      // sublane 0: [0xffffffff, 0xffffffff, ..., 0xffffffff]
-      // sublane 1: [0xffffffff, 0xffffffff, ..., 0xffffffff]
-      // sublane 2: [0xffffffff, 0xffffffff, ..., 0xffffffff]
-      // sublane 3: [0xffffffff, 0xffffffff, ..., 0xffffffff]
-      // sublane 4: [0xffffffff, 0xffffffff, ..., 0xffffffff]
-      // sublane 5: [0x0000ffff, 0x0000ffff, ..., 0x0000ffff]
-      // sublane 6: [0         , 0         , ..., 0         ]
-      // sublane 7: [0         , 0         , ..., 0         ]
-      //
-      // Through this way, in order to mask sub-elements, each target vreg only
-      // needs to apply 1 op (logical_and) instead of 3 ops (unpacking + select
-      // + packing).
-      Value partial_sublane_mask = builder.create<arith::ConstantOp>(
-          op.getLoc(),
-          DenseElementsAttr::get(
-              i32_vreg_ty,
-              builder.getI32IntegerAttr(
-                  0xffffffff >>
-                  (sub_padding * vreg_ty.getElementTypeBitWidth()))));
-      // Insert 0xffffffff above the blended sublane.
-      Value sublane_mask = builder.create<arith::SelectOp>(
-          getX32VmaskByPaddingEnd(0, x32_padding_bottom + 1), i32_max_vreg,
-          partial_sublane_mask);
-      // Insert 0 below the blended sublane.
-      sublane_mask = builder.create<arith::SelectOp>(mask_bottom, sublane_mask,
-                                                     i32_zeros_vreg);
-      for (int64_t i = 0; i < vregs.dim(1); ++i) {
-        Value &vreg = vregs({vregs.dim(0) - 1, i});
-        Value i32_vreg = builder.create<tpu::BitcastVregOp>(i32_vreg_ty, vreg);
-        if (sub_padding > 0) {
-          i32_vreg = builder.create<arith::AndIOp>(i32_vreg, sublane_mask);
-        } else {
-          i32_vreg = builder.create<arith::SelectOp>(mask_bottom, i32_vreg,
-                                                     i32_zeros_vreg);
-        }
-        vreg = builder.create<tpu::BitcastVregOp>(vreg_ty, i32_vreg);
-      }
-    }
-    // Mask out the right.
-    if (padding_right > 0) {
-      auto mask_right = getX32VmaskByPaddingEnd(1, padding_right);
-      for (int64_t i = 0; i < vregs.dim(0); ++i) {
-        Value &vreg = vregs({i, vregs.dim(1) - 1});
-        Value i32_vreg = builder.create<tpu::BitcastVregOp>(i32_vreg_ty, vreg);
-        i32_vreg = builder.create<arith::SelectOp>(mask_right, i32_vreg,
-                                                   i32_zeros_vreg);
-        vreg = builder.create<tpu::BitcastVregOp>(vreg_ty, i32_vreg);
-      }
-    }
-  };
-
-  // Create a vreg filled with zeros.
-  auto getZerosVergLike =
-      [&](const Value &vreg) -> FailureOr<TypedValue<VectorType>> {
-    const VectorType vreg_type = cast<VectorType>(vreg.getType());
-    FAILUREOR_ASSIGN_OR_RETURN(
-        const Attribute zero_attr,
-        getZeroIntOrFloatAttr(vreg_type.getElementType()));
-    return cast<TypedValue<VectorType>>(
-        builder
-            .create<arith::ConstantOp>(
-                op.getLoc(), DenseElementsAttr::get(vreg_type, zero_attr))
-            .getResult());
-  };
-
-  FAILUREOR_ASSIGN_OR_RETURN(auto lhs_zeros_vreg,
-                             getZerosVergLike(*lhs_vregs.begin()));
-  FAILUREOR_ASSIGN_OR_RETURN(auto rhs_zeros_vreg,
-                             getZerosVergLike(*rhs_vregs.begin()));
-  FAILUREOR_ASSIGN_OR_RETURN(auto acc_zeros_vreg,
-                             getZerosVergLike(*acc_vregs.begin()));
+  auto lhs_zeros_vreg =
+      getZerosVector(builder, cast<VectorType>(lhs_vregs.begin()->getType()));
+  auto rhs_zeros_vreg =
+      getZerosVector(builder, cast<VectorType>(rhs_vregs.begin()->getType()));
+  auto acc_zeros_vreg =
+      getZerosVector(builder, cast<VectorType>(acc_vregs.begin()->getType()));
 
   // Only mask out the paddings on contracting dim of LHS and RHS.
-  maskVregs(lhs_vregs, 0, padded_lhs_cols - lhs_shape[1]);
+  RETURN_IF_FAILED(
+      maskNativeTilingVregs(builder, lhs_vregs, ctx.target_shape,
+                            /*padding_bottom=*/0,
+                            /*padding_right=*/padded_lhs_cols - lhs_shape[1]));
   if (transpose_rhs) {
-    maskVregs(rhs_vregs, 0, padded_rhs_cols - rhs_shape[1]);
+    RETURN_IF_FAILED(maskNativeTilingVregs(
+        builder, rhs_vregs, ctx.target_shape,
+        /*padding_bottom=*/0,
+        /*padding_right=*/padded_rhs_cols - rhs_shape[1]));
   } else {
-    maskVregs(rhs_vregs, padded_rhs_rows - rhs_shape[0], 0);
+    RETURN_IF_FAILED(
+        maskNativeTilingVregs(builder, rhs_vregs, ctx.target_shape,
+                              /*padding_bottom=*/padded_rhs_rows - rhs_shape[0],
+                              /*padding_right=*/0));
   }
 
   // At this point, all paddings on vregs are masked out. For now, we
@@ -2875,12 +2739,10 @@ LogicalResult tpu_iota_rule(RewriteContext &ctx, Operation &op,
         native_vreg_ty,
         /*dimension =*/builder.getI32IntegerAttr(1));
     for (int64_t i = 0; i < num_tiles; ++i) {
-      auto offset = builder.create<arith::ConstantOp>(
-          native_vreg_ty,
-          DenseElementsAttr::get(
-              native_vreg_ty,
-              IntegerAttr::get(vty.getElementType(),
-                               i * *(native_vreg_ty.getShape().end() - 1))));
+      Value offset = getFullVector(
+          builder, native_vreg_ty,
+          IntegerAttr::get(vty.getElementType(),
+                           i * *(native_vreg_ty.getShape().end() - 1)));
       tiles[i] = builder.create<arith::AddIOp>(vreg_iota, offset);
     }
     xla::Array<Value> broadcasted_tiles(tile_array_shape);
@@ -2902,12 +2764,10 @@ LogicalResult tpu_iota_rule(RewriteContext &ctx, Operation &op,
         native_vreg_ty,
         /*dimension =*/builder.getI32IntegerAttr(0));
     for (int64_t i = 0; i < num_tiles; ++i) {
-      auto offset = builder.create<arith::ConstantOp>(
-          native_vreg_ty,
-          DenseElementsAttr::get(
-              native_vreg_ty,
-              IntegerAttr::get(vty.getElementType(),
-                               i * *(native_vreg_ty.getShape().end() - 2))));
+      Value offset = getFullVector(
+          builder, native_vreg_ty,
+          IntegerAttr::get(vty.getElementType(),
+                           i * *(native_vreg_ty.getShape().end() - 2)));
       tiles[i] = builder.create<arith::AddIOp>(vreg_iota, offset);
     }
     xla::Array<Value> broadcasted_tiles(tile_array_shape);
@@ -2924,10 +2784,8 @@ LogicalResult tpu_iota_rule(RewriteContext &ctx, Operation &op,
   SmallVector<Value> tiles;
   tiles.reserve(vty.getDimSize(*dimension));
   for (int64_t i = 0; i < vty.getDimSize(*dimension); ++i) {
-    tiles.push_back(builder.create<arith::ConstantOp>(
-        native_vreg_ty,
-        DenseElementsAttr::get(native_vreg_ty,
-                               IntegerAttr::get(vty.getElementType(), i))));
+    tiles.push_back(getFullVector(builder, native_vreg_ty,
+                                  IntegerAttr::get(vty.getElementType(), i)));
   }
   xla::Array<Value> out_tiles(tile_array_shape);
   out_tiles.Each([&](absl::Span<const int64_t> idxs, Value *v) {
@@ -3516,12 +3374,9 @@ LogicalResult vector_broadcast_rule(RewriteContext &ctx, Operation &op,
         const int64_t offset = *offsets_in[1];
         const int64_t lane_offset = offset % ctx.target_shape[1];
         const int64_t tile_offset = offset / ctx.target_shape[1];
-        const auto idx_ty =
-            VectorType::get(ctx.target_shape, builder.getI32Type());
-        auto lane_offset_cst = builder.create<arith::ConstantOp>(
-            broadcast_op.getLoc(), idx_ty,
-            DenseElementsAttr::get(idx_ty,
-                                   builder.getI32IntegerAttr(lane_offset)));
+        Value lane_offset_cst = getFullVector(
+            builder, getNativeVregType(builder.getI32Type(), ctx.target_shape),
+            builder.getI32IntegerAttr(lane_offset));
         DenseI32ArrayAttr sublane_pattern;
         if (num_tiles != 1) {
           SmallVector<int32_t> pattern;
@@ -3581,10 +3436,7 @@ LogicalResult vector_broadcast_rule(RewriteContext &ctx, Operation &op,
         getNativeVregType(src_i32.getType(), ctx.target_shape);
     auto tile_i32 =
         builder.create<vector::BroadcastOp>(native_vreg_ty, src_i32);
-    auto zeros = builder.create<arith::ConstantOp>(
-        broadcast_op.getLoc(), tile_i32.getType(),
-        DenseElementsAttr::get(tile_i32.getType(),
-                               builder.getI32IntegerAttr(0)));
+    Value zeros = getZerosVector(builder, tile_i32.getType());
     auto tile =
         builder.create<arith::CmpIOp>(arith::CmpIPredicate::ne, tile_i32, zeros)
             .getResult();
@@ -5479,8 +5331,6 @@ FailureOr<xla::Array<Value>> doColumnShiftRelayout(
   const std::array<int64_t, 2> vreg_slice = src.vregSlice(target_shape);
   const int bitwidth = src.bitwidth();
   const int packing = src.packing();
-  const VectorLayout dst(bitwidth, {src.offsets()[0], dst_col_offset}, tiling,
-                         src.implicit_dim());
   const int64_t col_diff = dst_col_offset - *src.offsets()[1];
   if (tiling[0] % packing != 0 || tiling[1] != target_shape[1]) {
     return emitError(loc,
@@ -5823,7 +5673,8 @@ LogicalResult retileToLargeTileWithScratch(
     RewriteContext &ctx, OpBuilder &builder, const Location loc,
     xla::Array<Value> &dst_tiles, const std::array<int64_t, 2> &dst_tile,
     const xla::Array<Value> &src_tiles, const std::array<int64_t, 2> &src_tile,
-    TypedValue<MemRefType> scratch_ref) {
+    TypedValue<MemRefType> scratch_ref, const int64_t store_vreg_delay,
+    const int64_t load_vreg_skips) {
   if (dst_tile[0] % src_tile[0] != 0) {
     return failure();
   }
@@ -5927,8 +5778,8 @@ LogicalResult retileToLargeTileWithScratch(
   SmallVector<int64_t, 4> src_idx(rank);
   dst_tiles.Each([&](absl::Span<const int64_t> dst_idx, Value *dst_vreg) {
     int64_t dst_row_idx = *(dst_idx.end() - 2);
-    int64_t dst_col_idx = *(dst_idx.end() - 1);
-    int64_t vreg_idx_in_group = dst_col_idx % vregs_per_group;
+    int64_t dst_col_idx_with_skips = *(dst_idx.end() - 1) + load_vreg_skips;
+    int64_t vreg_idx_in_group = dst_col_idx_with_skips % vregs_per_group;
     int64_t load_offset = sublanes_per_group * stored_group_cnt +
                           vreg_idx_in_group * sl_per_vreg * stride;
     delayed_loads.push_back(
@@ -5938,16 +5789,20 @@ LogicalResult retileToLargeTileWithScratch(
     // the vregs from current group and now we need to store corresponding
     // group of src vregs before actually emitting the loads.
     if (vreg_idx_in_group == vregs_per_group - 1 ||
-        dst_col_idx == dst_tiles.dimensions().back() - 1) {
-      auto src_row_idx = dst_row_idx * vregs_per_group;
-      auto src_col_idx = dst_col_idx / vregs_per_group;
+        dst_idx.back() == dst_tiles.dimensions().back() - 1) {
+      auto base_src_row_idx = dst_row_idx * vregs_per_group - store_vreg_delay;
+      auto src_col_idx = dst_col_idx_with_skips / vregs_per_group;
       std::copy(dst_idx.begin(), dst_idx.end(), src_idx.begin());
       for (int vi = 0; vi < vregs_per_group; ++vi) {
-        if (src_row_idx + vi >= src_tiles.dim(rank - 2) ||
+        const int64_t src_row_idx = base_src_row_idx + vi;
+        if (src_row_idx < 0) {
+          continue;
+        }
+        if (src_row_idx >= src_tiles.dim(rank - 2) ||
             src_col_idx >= src_tiles.dim(rank - 1)) {
           break;
         }
-        *(src_idx.end() - 2) = src_row_idx + vi;
+        *(src_idx.end() - 2) = src_row_idx;
         *(src_idx.end() - 1) = src_col_idx;
         Value src_vreg = src_tiles(src_idx);
         src_vreg =
@@ -5976,7 +5831,8 @@ LogicalResult retileToSmallTileWithScratch(
     RewriteContext &ctx, OpBuilder &builder, const Location loc,
     xla::Array<Value> &dst_tiles, const std::array<int64_t, 2> &dst_tile,
     const xla::Array<Value> &src_tiles, const std::array<int64_t, 2> &src_tile,
-    TypedValue<MemRefType> scratch_ref) {
+    TypedValue<MemRefType> scratch_ref, const int64_t store_vreg_delay,
+    const int64_t load_vreg_skips) {
   if (src_tile[0] % dst_tile[0] != 0) {
     return failure();
   }
@@ -6103,8 +5959,8 @@ LogicalResult retileToSmallTileWithScratch(
   SmallVector<int64_t, 4> dst_idx(rank);
   src_tiles.Each([&](absl::Span<const int64_t> src_idx, Value src_vreg) {
     int64_t src_row_idx = *(src_idx.end() - 2);
-    int64_t src_col_idx = *(src_idx.end() - 1);
-    int64_t vreg_idx_in_group = src_col_idx % vregs_per_group;
+    int64_t src_col_idx_with_delays = *(src_idx.end() - 1) + store_vreg_delay;
+    int64_t vreg_idx_in_group = src_col_idx_with_delays % vregs_per_group;
     src_vreg = builder.create<tpu::BitcastVregOp>(loc, temp_vreg_ty, src_vreg);
     if (use_shuffled_load) {
       Value store_offset = mlirIndexConst(
@@ -6126,16 +5982,20 @@ LogicalResult retileToSmallTileWithScratch(
     // vregs' row, this indicates we have stored all the vregs needed to
     // assemble a new group of dst vreg.
     if (vreg_idx_in_group == vregs_per_group - 1 ||
-        src_col_idx == src_tiles.dimensions().back() - 1) {
-      auto dst_row_idx = src_row_idx * vregs_per_group;
-      auto dst_col_idx = src_col_idx / vregs_per_group;
+        src_idx.back() == src_tiles.dimensions().back() - 1) {
+      auto base_dst_row_idx = src_row_idx * vregs_per_group - load_vreg_skips;
+      auto dst_col_idx = src_col_idx_with_delays / vregs_per_group;
       std::copy(src_idx.begin(), src_idx.end(), dst_idx.begin());
       for (int vi = 0; vi < vregs_per_group; ++vi) {
-        if (dst_row_idx + vi >= dst_tiles.dim(rank - 2) ||
+        const int64_t dst_row_idx = base_dst_row_idx + vi;
+        if (dst_row_idx < 0) {
+          continue;
+        }
+        if (dst_row_idx >= dst_tiles.dim(rank - 2) ||
             dst_col_idx >= dst_tiles.dim(rank - 1)) {
           break;
         }
-        *(dst_idx.end() - 2) = dst_row_idx + vi;
+        *(dst_idx.end() - 2) = dst_row_idx;
         *(dst_idx.end() - 1) = dst_col_idx;
         Value *dst_vreg = &dst_tiles(dst_idx);
         int64_t load_offset =
@@ -6160,18 +6020,70 @@ LogicalResult retileToSmallTileWithScratch(
 
 // go/mosaic-retiling-in-scratch is the full internal documentation that
 // includes more details about the TPU generations.
-LogicalResult retileWithScratch(RewriteContext &ctx, OpBuilder &builder,
-                                const Location loc,
-                                xla::Array<Value> &dst_tiles,
-                                const std::array<int64_t, 2> &dst_tiling,
-                                const xla::Array<Value> &src_tiles,
-                                const std::array<int64_t, 2> &src_tiling,
-                                int packing) {
+// Arguments:
+// - shape:            The non-implicit shape of the operand
+// - dst_tiling:       The desired result tiling
+// - dst_offsets_hint: Hints for the result offsets. They may be used or
+//                     ignored. See comments in the body of the function for
+//                     more details.
+// - src_vregs:        The source vregs to retile.
+// - src:              The source layout
+// Returns a pair holding the result layout (potentially using the hints) and
+// the retiled vregs.
+// TODO(tlongeri): Clean up the function parameters/signatures. We are passing
+// in more information than strictly needed.
+FailureOr<std::pair<VectorLayout, xla::Array<Value>>> retileWithScratch(
+    RewriteContext &ctx, OpBuilder &builder, const Location loc,
+    const ArrayRef<int64_t> shape, const std::array<int64_t, 2> dst_tiling,
+    const LayoutOffsets dst_offsets_hint, const xla::Array<Value> &src_vregs,
+    const VectorLayout &src) {
+  const int bitwidth = src.bitwidth();
+  const int packing = src.packing();
+  const std::array<int64_t, 2> src_tiling = src.tiling();
   if (!(src_tiling[1] == ctx.target_shape[1] &&
         dst_tiling[1] == ctx.target_shape[1] && src_tiling[0] % packing == 0 &&
         dst_tiling[0] % packing == 0)) {
     return failure();
   }
+  const std::array<int64_t, 2> src_vreg_slice =
+      VectorLayout::vregSlice(ctx.target_shape, bitwidth, src_tiling);
+  const std::array<int64_t, 2> dst_vreg_slice =
+      VectorLayout::vregSlice(ctx.target_shape, bitwidth, dst_tiling);
+
+  // TODO(b/368088671): When sublane tiling changes, we should be able to
+  // preserve some replications from the source layout. But we need to
+  // make sure they are implemented efficiently and well-tested. For now, we
+  // just simply use 0 for the replicated offset after retiling.
+  const LayoutOffsets src_offsets = {src.offsets()[0].value_or(0),
+                                     src.offsets()[1].value_or(0)};
+  // The provided offset hints are used only if they align with the source
+  // offsets, else we default to the smallest possible aligned offsets.
+  LayoutOffsets dst_offsets = {*src_offsets[0] % dst_vreg_slice[0],
+                               *src_offsets[1] % dst_vreg_slice[1]};
+  // On a given dimension, either the source vreg slice size divides the dest
+  // vreg slice size, or vice versa (depending on the dimension and whether it's
+  // small-to-large or large-to-small retiling). Offset changes are supported
+  // as long as they are aligned modulo the smaller of the two sizes.
+  const std::array<int64_t, 2> alignment = {
+      std::min(src_vreg_slice[0], dst_vreg_slice[0]),
+      std::min(src_vreg_slice[1], dst_vreg_slice[1])};
+  if (dst_offsets_hint[0].has_value() &&
+      (*dst_offsets_hint[0] - *src_offsets[0]) % alignment[0] == 0) {
+    CHECK_LT(*dst_offsets_hint[0], dst_vreg_slice[0]);
+    dst_offsets[0] = *dst_offsets_hint[0];
+  }
+  if (dst_offsets_hint[1].has_value() &&
+      (*dst_offsets_hint[1] - *src_offsets[1]) % alignment[1] == 0) {
+    CHECK_LT(*dst_offsets_hint[1], dst_vreg_slice[1]);
+    dst_offsets[1] = *dst_offsets_hint[1];
+  }
+  // The offsets of the source in units of the destination vreg slice:
+  const std::array<int64_t, 2> src_offsets_in_dst_vreg_slices = {
+      *src_offsets[0] / dst_vreg_slice[0], *src_offsets[1] / dst_vreg_slice[1]};
+  // The offsets of the destination in units of the source vreg slice:
+  const std::array<int64_t, 2> dst_offsets_in_src_vreg_slices = {
+      *dst_offsets[0] / src_vreg_slice[0], *dst_offsets[1] / src_vreg_slice[1]};
+
   // Try to get i32 vector scratch space. Because we will bitcast vregs to
   // i32 vregs before using scratch for retiling. Through this way we can
   // handle packed types as well.
@@ -6186,24 +6098,57 @@ LogicalResult retileWithScratch(RewriteContext &ctx, OpBuilder &builder,
                                             dst_tiling[1]};
   std::array<int64_t, 2> vi32_src_tiling = {src_tiling[0] / packing,
                                             src_tiling[1]};
+
+  const VectorLayout dst(bitwidth, dst_offsets, dst_tiling, src.implicit_dim());
+  TPU_ASSERT_LOC(loc, dst.isValid(ctx.target_shape));
+  xla::Array<Value> dst_vregs(
+      dst.tileArrayImplicitShape(shape, ctx.target_shape));
+  // When differences in offsets exist, the source vregs may stored at an offset
+  // position in their group. For example, the 1st vreg in a row/column may be
+  // stored as if it was the 3rd, so that the parts corresponding to the 1st and
+  // 2nd in the destination are filled with padding. Likewise, loads to
+  // destination vregs may be skipped, when they would load only padding.
+  // store_vreg_delay is the position offset for stores, and load_vreg_skips is
+  // the position offset for loads.
+  //
+  // For example, suppose we are going from 32-bit {0, 128}(2, 128) to
+  // {4, 0}(8, 128). We form groups of 4 vregs that represent an (8, 512) slice
+  // of the padded implicit shape. For the given offsets, for the first group,
+  // the data is in (4:8, 128:512). But the first and second sources (stored
+  // vregs) of the group form the slices of data (0:2, 0:512) and (2:4, 0:512),
+  // which should be all padding. Likewise, the first dest vreg slice (which we
+  // load from) holds the data from slice (0:8, 0:128), which is all padding.
+  // We never load or store to slices that should contain only padding.
   if (src_tiling[0] > dst_tiling[0]) {
-    return retileToSmallTileWithScratch(ctx, builder, loc, dst_tiles,
-                                        vi32_dst_tiling, src_tiles,
-                                        vi32_src_tiling, ref);
+    DCHECK_EQ(src_offsets_in_dst_vreg_slices[1], 0);
+    DCHECK_EQ(dst_offsets_in_src_vreg_slices[0], 0);
+    const int64_t store_vreg_delay = dst_offsets_in_src_vreg_slices[1];
+    const int64_t load_vreg_skips = src_offsets_in_dst_vreg_slices[0];
+    if (failed(retileToSmallTileWithScratch(
+            ctx, builder, loc, dst_vregs, vi32_dst_tiling, src_vregs,
+            vi32_src_tiling, ref, store_vreg_delay, load_vreg_skips))) {
+      return failure();
+    }
   }
   if (src_tiling[0] < dst_tiling[0]) {
-    return retileToLargeTileWithScratch(ctx, builder, loc, dst_tiles,
-                                        vi32_dst_tiling, src_tiles,
-                                        vi32_src_tiling, ref);
+    DCHECK_EQ(src_offsets_in_dst_vreg_slices[0], 0);
+    DCHECK_EQ(dst_offsets_in_src_vreg_slices[1], 0);
+    const int64_t store_vreg_delay = dst_offsets_in_src_vreg_slices[0];
+    const int64_t load_vreg_skips = src_offsets_in_dst_vreg_slices[1];
+    if (failed(retileToLargeTileWithScratch(
+            ctx, builder, loc, dst_vregs, vi32_dst_tiling, src_vregs,
+            vi32_src_tiling, ref, store_vreg_delay, load_vreg_skips))) {
+      return failure();
+    }
   }
-  dst_tiles = std::move(src_tiles);
-  return success();
+  return std::make_pair(dst, dst_vregs);
 }
 
 FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
     RewriteContext &ctx, OpBuilder &builder, const Location loc, VectorType vty,
     const VectorLayout src, xla::Array<Value> vregs,
-    const std::array<int64_t, 2> dst_tiling, bool try_replicate_rows) {
+    const std::array<int64_t, 2> dst_tiling,
+    const LayoutOffsets dst_offsets_hint) {
   bool has_enough_scratch = ctx.max_sublanes_in_scratch >=
                             ctx.target_shape[0] * (ctx.target_shape[0] + 1);
   const auto &target_shape = ctx.target_shape;
@@ -6219,6 +6164,12 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
   const int8_t bitwidth = src.bitwidth();
   const std::array<int64_t, 2> dst_vreg_slice =
       VectorLayout::vregSlice(ctx.target_shape, bitwidth, dst_tiling);
+  // TODO(tlongeri): Using canonical vs non-canonical offsets can change the
+  // value of try_replicate rows, and it breaks some tests. It doesn't make
+  // sense that we have different behavior for equivalent layouts, though. We
+  // need better logic for picking the relayout strategy.
+  const bool try_replicate_rows =
+      src.offsets()[0].has_value() && !dst_offsets_hint[0].has_value();
 
   // Fully replicated offsets are handled efficiently elsewhere (in relayout)
   CHECK(src.offsets()[0].has_value() || src.offsets()[1].has_value());
@@ -6290,15 +6241,10 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
     });
     return std::pair(dst, std::move(retiled));
   }
-  VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
-                   src.implicit_dim());
-  if (!dst.isValid(target_shape)) {
-    return emitError(loc, "Not implemented: invalid offsets in tiling target");
-  }
-  auto dst_tiles_shape =
-      dst.tileArrayImplicitShape(vty.getShape(), target_shape);
   // (8,128) -> (8 * packing,128) tiling change for packed type.
-  if (bitwidth < 32 && 32 % bitwidth == 0 && src_tiling == ctx.target_shape &&
+  if (src_offsets[0].value_or(0) < dst_vreg_slice[0] &&
+      src_offsets[1].value_or(0) < dst_vreg_slice[1] && bitwidth < 32 &&
+      32 % bitwidth == 0 && src_tiling == ctx.target_shape &&
       dst_tiling == std::array<int64_t, 2>{ctx.target_shape[0] * packing,
                                            ctx.target_shape[1]}) {
     // TODO(tlongeri): This relayout is just ext + trunc. Refactor.
@@ -6308,8 +6254,10 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
       // not, since it relies on the src vreg array shape to know how many tiles
       // to pack in dst, and vreg array shapes with materialized offsets are
       // unfortunately not equal to vreg array shapes with replicated offsets.
-      CHECK(dst.offsets() == src_offsets);
-      xla::Array<Value> retiled(dst_tiles_shape);
+      VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
+                       src.implicit_dim());
+      xla::Array<Value> retiled(
+          dst.tileArrayImplicitShape(vty.getShape(), target_shape));
       VectorType vreg_x32 =
           vty.getElementType().isSignlessInteger()
               ? VectorType::get(target_shape, builder.getI32Type())
@@ -6357,7 +6305,9 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
   // interesting if the next step is a retile, since we can also
   // match corresponding elements without shifting. It's just that
   // the tiles are not adjacent (no contiguous vreg slice).
-  if (bitwidth < 32 && 32 % bitwidth == 0 &&
+  if (src_offsets[0].value_or(0) < dst_vreg_slice[0] &&
+      src_offsets[1].value_or(0) < dst_vreg_slice[1] && bitwidth < 32 &&
+      32 % bitwidth == 0 &&
       src_tiling == std::array<int64_t, 2>{1, ctx.target_shape[1] * packing} &&
       dst_tiling == std::array<int64_t, 2>{packing, ctx.target_shape[1]}) {
     // TODO(tlongeri): This relayout is just ext + trunc. Refactor.
@@ -6406,8 +6356,10 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
     // not, since it relies on the src vreg array shape to know how many tiles
     // to pack in dst, and vreg array shapes with materialized offsets are
     // unfortunately not equal to vreg array shapes with replicated offsets.
-    CHECK(dst.offsets() == src.offsets());
-    xla::Array<Value> retiled(dst_tiles_shape);
+    VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
+                     src.implicit_dim());
+    xla::Array<Value> retiled(
+        dst.tileArrayImplicitShape(vty.getShape(), target_shape));
     const VectorType vreg_x32 =
         vty.getElementType().isSignlessInteger()
             ? VectorType::get(target_shape, builder.getI32Type())
@@ -6444,24 +6396,25 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
     return std::pair(dst, std::move(retiled));
   }
   if (src_tiling[1] == target_shape[1] && dst_tiling[1] == target_shape[1]) {
-    // TODO(b/368088671): When sublane tiling changes, we should be able to
-    // preserve some replications from the source layout. But we need to
-    // make sure they are implemented efficiently and well-tested. For now, we
-    // just simply use 0 for the replicated offset after retiling.
-    dst = VectorLayout(
-        bitwidth, {src.offsets()[0].value_or(0), src.offsets()[1].value_or(0)},
-        dst_tiling, dst.implicit_dim());
-
     // All clauses in the and expression are based on performance benchmarking.
     bool use_alu = !has_enough_scratch ||
                    (ctx.hardware_generation >= 5 && src_tiling[0] != packing &&
                     dst_tiling[0] != packing);
 
     if (use_alu) {
-      if (src_tiling[0] > dst_tiling[0]) {
-        return std::pair(
-            dst, retileToReducedSublanes(builder, vty.getShape(), src, vregs,
-                                         dst, target_shape));
+      if (src_tiling[0] > dst_tiling[0] &&
+          // retileToReducedSublanes does not support offset changes
+          src.offsets()[0].value_or(0) < dst_vreg_slice[0] &&
+          src.offsets()[1].value_or(0) < dst_vreg_slice[1]) {
+        VectorLayout dst(src.bitwidth(), src.offsets(), dst_tiling,
+                         src.implicit_dim());
+        return std::pair(dst, retileToReducedSublanes(
+                                  builder, vty.getShape(), src, vregs,
+                                  VectorLayout(bitwidth,
+                                               {src.offsets()[0].value_or(0),
+                                                src.offsets()[1].value_or(0)},
+                                               dst_tiling, dst.implicit_dim()),
+                                  target_shape));
       } else if (!has_enough_scratch) {
         // TODO(b/357538782): Implement retileToIncreasedSublanes with ALU ops.
         return emitError(
@@ -6469,15 +6422,12 @@ FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeTiling(
             "Not implemented: retiling to increase sublane tiling with ALU");
       }
     }
-    xla::Array<Value> retiled(dst_tiles_shape);
-    if (failed(retileWithScratch(ctx, builder, loc, retiled, dst_tiling, vregs,
-                                 src_tiling, packing))) {
-      return failure();
-    }
-    return std::pair(dst, std::move(retiled));
+    return retileWithScratch(ctx, builder, loc, vty.getShape(), dst_tiling,
+                             dst_offsets_hint, vregs, src);
   }
   return emitError(loc, "Not implemented: Unsupported tiling change for ")
-         << vty << ": from " << src << " to " << dst;
+         << vty << ": from " << src << " to (" << dst_tiling[0] << ", "
+         << dst_tiling[1] << ") tiling";
 }
 
 FailureOr<std::pair<VectorLayout, xla::Array<Value>>> changeImplicitDim(
@@ -6737,9 +6687,7 @@ FailureOr<TypedValue<VectorType>> relayout(RewriteContext &ctx,
   FAILUREOR_ASSIGN_OR_RETURN(
       std::tie(src, src_tiles),
       changeTiling(ctx, builder, v.getLoc(), vty, src, std::move(src_tiles),
-                   dst.tiling(),
-                   dst.offsets()[0] == std::nullopt &&
-                       src.offsets()[0] != std::nullopt));
+                   dst.tiling(), dst.offsets()));
 
   FAILUREOR_ASSIGN_OR_RETURN(
       std::tie(src, src_tiles),
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
index 43466c82ac89..c5448a3df514 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -1647,10 +1647,17 @@ class VectorLayoutInferer {
     Layout dst_layout;
     if (layout.tiling() == nativeTiling(src_bitwidth)) {
       // If the source is already in native tiling, we can unpack it directly.
-      src_layout = layout;
+      std::array<int64_t, 2> dst_native_tiling = nativeTiling(dst_bitwidth);
+      LayoutOffsets offsets = {layout.offsets()[0]
+                                   ? *layout.offsets()[0] % dst_native_tiling[0]
+                                   : LayoutOffset(),
+                               layout.offsets()[1]};
+      DCHECK_LT(offsets[1].value_or(0), dst_native_tiling[1]);
+      src_layout = VectorLayout(src_bitwidth, offsets, layout.tiling(),
+                                layout.implicit_dim());
       dst_layout =
-          VectorLayout(dst_bitwidth, layout.offsets(),
-                       nativeTiling(dst_bitwidth), layout.implicit_dim());
+          VectorLayout(dst_bitwidth, offsets, dst_native_tiling,
+                       layout.implicit_dim());
     } else if (dst_bitwidth == 32 &&
                default_tiling_[0] % layout.tiling()[0] == 0 &&
                default_tiling_[1] == layout.tiling()[1]) {
@@ -1659,13 +1666,17 @@ class VectorLayoutInferer {
       // tiling through the op.
       // TODO(jevinjiang): we can relax this for non-32bit as well.
       src_layout = layout;
-      dst_layout = VectorLayout(32, layout.offsets(), src_layout->tiling(),
-                                layout.implicit_dim());
+      dst_layout = VectorLayout(dst_bitwidth, layout.offsets(),
+                                src_layout->tiling(), layout.implicit_dim());
     } else {
-      // TODO(b/335863273): we should also reduce offsets.
-      src_layout = VectorLayout(src_bitwidth, layout.offsets(), default_tiling_,
+      LayoutOffsets offsets = {
+          layout.offsets()[0] ? *layout.offsets()[0] % default_tiling_[0]
+                              : LayoutOffset(),
+          layout.offsets()[1] ? *layout.offsets()[1] % default_tiling_[1]
+                              : LayoutOffset()};
+      src_layout = VectorLayout(src_bitwidth, offsets, default_tiling_,
                                 layout.implicit_dim());
-      dst_layout = VectorLayout(dst_bitwidth, layout.offsets(), default_tiling_,
+      dst_layout = VectorLayout(dst_bitwidth, offsets, default_tiling_,
                                 layout.implicit_dim());
     }
     setLayout(op, src_layout, dst_layout);
diff --git a/jaxlib/mosaic/dialect/tpu/vreg_util.cc b/jaxlib/mosaic/dialect/tpu/vreg_util.cc
new file mode 100644
index 000000000000..7dc5c13c073e
--- /dev/null
+++ b/jaxlib/mosaic/dialect/tpu/vreg_util.cc
@@ -0,0 +1,206 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "jaxlib/mosaic/dialect/tpu/vreg_util.h"
+
+#include <array>
+#include <cstdint>
+
+#include "absl/log/check.h"
+#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/include/mlir/IR/Attributes.h"
+#include "mlir/include/mlir/IR/BuiltinAttributes.h"
+#include "mlir/include/mlir/IR/BuiltinTypes.h"
+#include "mlir/include/mlir/IR/Diagnostics.h"
+#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/include/mlir/IR/Types.h"
+#include "mlir/include/mlir/IR/Value.h"
+#include "mlir/include/mlir/Support/LLVM.h"
+#include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
+#include "jaxlib/mosaic/dialect/tpu/util.h"
+#include "xla/array.h"
+
+namespace mlir::tpu {
+
+namespace {
+
+VectorType getNativeVregOrVmaskTypeImpl(
+    Type elem_ty, const int8_t bitwidth,
+    const std::array<int64_t, 2> target_shape) {
+  if (bitwidth == 32) {
+    return VectorType::get(target_shape, elem_ty);
+  }
+  return VectorType::get({target_shape[0], target_shape[1], 32 / bitwidth},
+                         elem_ty);
+}
+
+}  // namespace
+
+VectorType getNativeVregOrVmaskType(Type elem_ty, const int8_t layout_bitwidth,
+                                    const std::array<int64_t, 2> target_shape) {
+  int8_t bitwidth = elem_ty.getIntOrFloatBitWidth();
+  if (bitwidth == 1) {
+    bitwidth = layout_bitwidth;
+  } else {
+    CHECK_EQ(bitwidth, layout_bitwidth);
+  }
+  return getNativeVregOrVmaskTypeImpl(elem_ty, bitwidth, target_shape);
+}
+
+VectorType getNativeVregType(Type elem_ty,
+                             const std::array<int64_t, 2> target_shape) {
+  return getNativeVregOrVmaskTypeImpl(elem_ty, elem_ty.getIntOrFloatBitWidth(),
+                                      target_shape);
+}
+
+TypedValue<VectorType> getFullVector(ImplicitLocOpBuilder &builder,
+                                     VectorType vty, Attribute value) {
+  return cast<TypedValue<VectorType>>(
+      builder.create<arith::ConstantOp>(DenseElementsAttr::get(vty, value))
+          .getResult());
+}
+
+TypedValue<VectorType> getFullLikeVector(ImplicitLocOpBuilder &builder,
+                                         TypedValue<VectorType> vec,
+                                         Attribute value) {
+  return getFullVector(builder, vec.getType(), value);
+}
+
+TypedValue<VectorType> getZerosVector(ImplicitLocOpBuilder &builder,
+                                      VectorType vty) {
+  return getFullVector(builder, vty, builder.getZeroAttr(vty.getElementType()));
+}
+
+TypedValue<VectorType> getZerosLikeVector(ImplicitLocOpBuilder &builder,
+                                          TypedValue<VectorType> vec) {
+  return getZerosVector(builder, vec.getType());
+}
+
+FailureOr<TypedValue<VectorType>> getX32VmaskByPaddingEnd(
+    ImplicitLocOpBuilder &builder, int64_t padding,
+    const std::array<int64_t, 2> target_shape, int64_t dim) {
+  VectorType i32_vreg_ty =
+      getNativeVregType(builder.getI32Type(), target_shape);
+  if (dim != 0 && dim != 1) {
+    return builder.emitError()
+           << "Expected a 2D vector for getX32VmaskByPaddingEnd";
+  }
+
+  if (padding < 0 || padding > target_shape[dim]) {
+    return builder.emitError()
+           << "Padding must be in [0, target_shape[dim]). Padding: " << padding
+           << ", target_shape[dim]: " << target_shape[dim];
+  }
+
+  Value padding_vreg =
+      getFullVector(builder, i32_vreg_ty,
+                    builder.getI32IntegerAttr(target_shape[dim] - padding));
+
+  return cast<TypedValue<VectorType>>(
+      builder
+          .create<arith::CmpIOp>(
+              arith::CmpIPredicate::slt,
+              builder.create<tpu::IotaOp>(i32_vreg_ty,
+                                          builder.getI32IntegerAttr(dim)),
+              padding_vreg)
+          .getResult());
+}
+
+LogicalResult maskNativeTilingVregs(ImplicitLocOpBuilder &builder,
+                                    xla::Array<Value> &vregs,
+                                    std::array<int64_t, 2> target_shape,
+                                    int64_t padding_bottom,
+                                    int64_t padding_right) {
+  auto vreg_ty = dyn_cast<VectorType>(vregs.begin()->getType());
+  if (!vreg_ty) {
+    return builder.emitError() << "Expected a vector type";
+  }
+
+  VectorType i32_vreg_ty =
+      getNativeVregType(builder.getI32Type(), target_shape);
+  Value i32_zeros_vreg = getZerosVector(builder, i32_vreg_ty);
+  Value i32_max_vreg = getFullVector(builder, i32_vreg_ty,
+                                     builder.getI32IntegerAttr(0xffffffff));
+
+  int packing = vreg_ty.getRank() > 2 ? vreg_ty.getShape()[2] : 1;
+  // Mask out the bottom.
+  if (padding_bottom > 0) {
+    // The function is only called when the vreg has native tiling. Therefore,
+    // it is safe to bitcast to x32 vreg for masking.
+    int sub_padding = padding_bottom % packing;
+    int x32_padding_bottom = padding_bottom / packing;
+    FAILUREOR_ASSIGN_OR_RETURN(
+        Value mask_top, getX32VmaskByPaddingEnd(builder, x32_padding_bottom + 1,
+                                                target_shape, /*dim=*/0));
+    FAILUREOR_ASSIGN_OR_RETURN(
+        Value mask_bottom,
+        getX32VmaskByPaddingEnd(builder, x32_padding_bottom, target_shape,
+                                /*dim=*/0));
+    // Create an int32 vreg which contains subelement masking and then
+    // logical_and with target vreg to mask out the unaligned paddings.
+    // Eg. if padding_bottom = 5, packing = 2, and assume the vreg shape is
+    // [8, 128], then the mask will be:
+    //
+    // sublane 0: [0xffffffff, 0xffffffff, ..., 0xffffffff]
+    // sublane 1: [0xffffffff, 0xffffffff, ..., 0xffffffff]
+    // sublane 2: [0xffffffff, 0xffffffff, ..., 0xffffffff]
+    // sublane 3: [0xffffffff, 0xffffffff, ..., 0xffffffff]
+    // sublane 4: [0xffffffff, 0xffffffff, ..., 0xffffffff]
+    // sublane 5: [0x0000ffff, 0x0000ffff, ..., 0x0000ffff]
+    // sublane 6: [0         , 0         , ..., 0         ]
+    // sublane 7: [0         , 0         , ..., 0         ]
+    //
+    // Through this way, in order to mask sub-elements, each target vreg only
+    // needs to apply 1 op (logical_and) instead of 3 ops (unpacking + select
+    // + packing).
+    Value partial_sublane_mask = getFullVector(
+        builder, i32_vreg_ty,
+        builder.getI32IntegerAttr(
+            0xffffffff >> (sub_padding * vreg_ty.getElementTypeBitWidth())));
+    // Insert 0xffffffff above the blended sublane.
+    Value sublane_mask = builder.create<arith::SelectOp>(mask_top, i32_max_vreg,
+                                                         partial_sublane_mask);
+    // Insert 0 below the blended sublane.
+    sublane_mask = builder.create<arith::SelectOp>(mask_bottom, sublane_mask,
+                                                   i32_zeros_vreg);
+    for (int64_t i = 0; i < vregs.dim(1); ++i) {
+      Value &vreg = vregs({vregs.dim(0) - 1, i});
+      Value i32_vreg = builder.create<tpu::BitcastVregOp>(i32_vreg_ty, vreg);
+      if (sub_padding > 0) {
+        i32_vreg = builder.create<arith::AndIOp>(i32_vreg, sublane_mask);
+      } else {
+        i32_vreg = builder.create<arith::SelectOp>(mask_bottom, i32_vreg,
+                                                   i32_zeros_vreg);
+      }
+      vreg = builder.create<tpu::BitcastVregOp>(vreg_ty, i32_vreg);
+    }
+  }
+  // Mask out the right.
+  if (padding_right > 0) {
+    FAILUREOR_ASSIGN_OR_RETURN(
+        Value mask_right, getX32VmaskByPaddingEnd(builder, padding_right,
+                                                  target_shape, /*dim=*/1));
+    for (int64_t i = 0; i < vregs.dim(0); ++i) {
+      Value &vreg = vregs({i, vregs.dim(1) - 1});
+      Value i32_vreg = builder.create<tpu::BitcastVregOp>(i32_vreg_ty, vreg);
+      i32_vreg =
+          builder.create<arith::SelectOp>(mask_right, i32_vreg, i32_zeros_vreg);
+      vreg = builder.create<tpu::BitcastVregOp>(vreg_ty, i32_vreg);
+    }
+  }
+  return success();
+}
+
+}  // namespace mlir::tpu
diff --git a/jaxlib/mosaic/dialect/tpu/vreg_util.h b/jaxlib/mosaic/dialect/tpu/vreg_util.h
new file mode 100644
index 000000000000..5892582a9f4a
--- /dev/null
+++ b/jaxlib/mosaic/dialect/tpu/vreg_util.h
@@ -0,0 +1,82 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_PY_JAX_JAXLIB_MOSAIC_DIALECT_TPU_VREG_UTIL_H_
+#define THIRD_PARTY_PY_JAX_JAXLIB_MOSAIC_DIALECT_TPU_VREG_UTIL_H_
+
+#include <array>
+#include <cstdint>
+
+#include "mlir/include/mlir/IR/Attributes.h"
+#include "mlir/include/mlir/IR/Builders.h"
+#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/include/mlir/IR/Types.h"
+#include "mlir/include/mlir/IR/Value.h"
+#include "mlir/include/mlir/Support/LLVM.h"
+#include "xla/array.h"
+
+namespace mlir::tpu {
+
+// Returns the native vreg or vmask type for the given element type and target
+// shape. The layout bitwidth is used for i1 (vmask) elements.
+VectorType getNativeVregOrVmaskType(Type elem_ty, int8_t layout_bitwidth,
+                                    std::array<int64_t, 2> target_shape);
+VectorType getNativeVregType(Type elem_ty, std::array<int64_t, 2> target_shape);
+
+// Returns a zero constant of the same type as `vty`.
+TypedValue<VectorType> getZerosVector(ImplicitLocOpBuilder &builder,
+                                      VectorType vty);
+// Same as above, but takes a `vec` as input.
+TypedValue<VectorType> getZerosLikeVector(ImplicitLocOpBuilder &builder,
+                                          TypedValue<VectorType> vec);
+
+// Returns a constant of the same type as `vty` with the given `value`.
+TypedValue<VectorType> getFullVector(ImplicitLocOpBuilder &builder,
+                                     VectorType vty, Attribute value);
+// Same as above, but takes a `vec` as input.
+TypedValue<VectorType> getFullLikeVector(ImplicitLocOpBuilder &builder,
+                                         TypedValue<VectorType> vec,
+                                         Attribute value);
+
+// Creates a vmask with false flags to bottom (dim = 0)
+// or right (dim = 1) where the flag count corresponds to the (dim_size -
+// padding).
+//
+// For example, assume vmask shape is (4, 8)
+//
+// getX32VmaskByPaddingEnd(padding=3, dim=1) creates:
+//  [T, T, T, T, T, F, F, F]
+//  [T, T, T, T, T, F, F, F]
+//  [T, T, T, T, T, F, F, F]
+//  [T, T, T, T, T, F, F, F]
+// TODO(b/385204135): Unify with getVmaskByPaddingEnd in tpu_rotate_rule, and
+// improve the codegen.
+FailureOr<TypedValue<VectorType>> getX32VmaskByPaddingEnd(
+    ImplicitLocOpBuilder &builder, int64_t padding,
+    std::array<int64_t, 2> target_shape, int64_t dim);
+
+// Masks out the padding in the bottom and right of the vregs. vregs are
+// expected to have native tiling, and the masked vregs are mutated in
+// `vregs`. `padding_bottom` and `padding_right` is the number of elements to
+// pad in the bottom and right.
+LogicalResult maskNativeTilingVregs(ImplicitLocOpBuilder &builder,
+                                    xla::Array<Value> &vregs,
+                                    std::array<int64_t, 2> target_shape,
+                                    int64_t padding_bottom,
+                                    int64_t padding_right);
+
+}  // namespace mlir::tpu
+
+#endif  // THIRD_PARTY_PY_JAX_JAXLIB_MOSAIC_DIALECT_TPU_VREG_UTIL_H_
diff --git a/jaxlib/mosaic/dialect/tpu/vreg_util_test.cc b/jaxlib/mosaic/dialect/tpu/vreg_util_test.cc
new file mode 100644
index 000000000000..dadbac133fbf
--- /dev/null
+++ b/jaxlib/mosaic/dialect/tpu/vreg_util_test.cc
@@ -0,0 +1,228 @@
+/* Copyright 2024 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "jaxlib/mosaic/dialect/tpu/vreg_util.h"
+
+#include <array>
+#include <cstdint>
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/include/mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/include/mlir/IR/Builders.h"
+#include "mlir/include/mlir/IR/BuiltinOps.h"
+#include "mlir/include/mlir/IR/BuiltinTypes.h"
+#include "mlir/include/mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/include/mlir/IR/MLIRContext.h"
+#include "mlir/include/mlir/IR/OwningOpRef.h"
+#include "mlir/include/mlir/IR/Value.h"
+#include "mlir/include/mlir/Support/DebugStringHelper.h"
+#include "mlir/include/mlir/Support/LLVM.h"
+#include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
+
+namespace mlir::tpu {
+
+namespace {
+
+using ::testing::Eq;
+using ::testing::Optional;
+
+MATCHER_P2(IsConstantOpWithSplatValue, type, splat_value, "") {
+  auto constant_op = dyn_cast<arith::ConstantOp>(arg.getDefiningOp());
+  if (constant_op == nullptr) {
+    *result_listener << "Expected a constant op, got " << debugString(arg);
+    return false;
+  }
+  auto dense_attr = dyn_cast<DenseElementsAttr>(constant_op.getValue());
+  if (dense_attr == nullptr) {
+    *result_listener << "Expected a dense elements attr, got "
+                     << debugString(arg);
+    return false;
+  }
+  if (dense_attr.getType() != type) {
+    *result_listener << "Expected a dense elements attr with type "
+                     << debugString(type) << ", got "
+                     << debugString(dense_attr.getType());
+    return false;
+  }
+  if (!dense_attr.isSplat()) {
+    *result_listener << "Expected a splat dense elements attr, got "
+                     << debugString(dense_attr);
+    return false;
+  }
+  if (auto s = dense_attr.template getSplatValue<decltype(splat_value)>();
+      s != splat_value) {
+    *result_listener << "Expected a splat dense elements attr with value "
+                     << splat_value << ", got " << s;
+    return false;
+  }
+  return true;
+}
+
+MATCHER_P2(IsVectorTypeWithShape, shape, elem_ty, "") {
+  auto vty = dyn_cast<VectorType>(arg);
+  if (vty == nullptr) {
+    *result_listener << "Expected a vector type, got " << debugString(arg);
+    return false;
+  }
+  if (vty.getShape() != ArrayRef<int64_t>(shape)) {
+    *result_listener << "Expected a vector type with shape "
+                     << absl::StrJoin(shape, ",") << ", got "
+                     << absl::StrJoin(vty.getShape(), ",");
+    return false;
+  }
+  if (vty.getElementType() != elem_ty) {
+    *result_listener << "Expected a vector type with element type "
+                     << debugString(elem_ty) << ", got "
+                     << debugString(vty.getElementType());
+    return false;
+  }
+  return true;
+}
+
+class VregUtilTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    context_.loadDialect<arith::ArithDialect, vector::VectorDialect,
+                         tpu::TPUDialect>();
+    mlir::Location loc = mlir::UnknownLoc::get(&context_);
+    mlir::OpBuilder b(&context_);
+    module_ = b.create<ModuleOp>(loc);
+    builder_ = std::make_unique<mlir::ImplicitLocOpBuilder>(
+        module_->getLoc(), module_->getBodyRegion());
+  }
+
+  void TearDown() override {
+    builder_.reset();
+    // Reset the module to prevent memory leaks.
+    module_ = nullptr;
+  }
+
+  mlir::ImplicitLocOpBuilder& Builder() { return *builder_; }
+
+ private:
+  MLIRContext context_;
+  std::unique_ptr<mlir::ImplicitLocOpBuilder> builder_;
+  OwningOpRef<ModuleOp> module_;
+};
+
+TEST_F(VregUtilTest, GetNativeVregOrVmaskTypeBitwidthMismatch) {
+  EXPECT_DEATH(getNativeVregOrVmaskType(Builder().getI16Type(),
+                                        /*layout_bitwidth=*/8, {2, 4}),
+               "");
+}
+
+TEST_F(VregUtilTest, GetNativeVregOrVmaskTypeI1) {
+  EXPECT_THAT(getNativeVregOrVmaskType(Builder().getI1Type(),
+                                       /*layout_bitwidth=*/8, {2, 4}),
+              IsVectorTypeWithShape(std::array<int64_t, 3>{2, 4, 4},
+                                    Builder().getI1Type()));
+}
+
+TEST_F(VregUtilTest, GetNativeVregF32) {
+  EXPECT_THAT(getNativeVregType(Builder().getF32Type(), {2, 4}),
+              IsVectorTypeWithShape(std::array<int64_t, 2>{2, 4},
+                                    Builder().getF32Type()));
+}
+
+TEST_F(VregUtilTest, GetNativeVregBf16) {
+  EXPECT_THAT(getNativeVregType(Builder().getBF16Type(), {2, 4}),
+              IsVectorTypeWithShape(std::array<int64_t, 3>{2, 4, 2},
+                                    Builder().getBF16Type()));
+}
+
+TEST_F(VregUtilTest, GetFullVector) {
+  VectorType vty = VectorType::get({2, 4}, Builder().getI32Type());
+  TypedValue<VectorType> vec =
+      getFullVector(Builder(), vty, Builder().getI32IntegerAttr(0x1));
+
+  EXPECT_THAT(vec, IsConstantOpWithSplatValue(vty, int32_t{0x1}));
+}
+
+TEST_F(VregUtilTest, GetFullLikeVector) {
+  VectorType vty = VectorType::get({2, 4}, Builder().getF32Type());
+  TypedValue<VectorType> in_vec = Builder().create<vector::SplatOp>(
+      vty, Builder().create<arith::ConstantOp>(
+               vty.getElementType(), Builder().getF32FloatAttr(1.0f)));
+  TypedValue<VectorType> vec =
+      getFullLikeVector(Builder(), in_vec, Builder().getF32FloatAttr(2.0f));
+
+  EXPECT_THAT(vec, IsConstantOpWithSplatValue(vty, float{2.0f}));
+}
+
+TEST_F(VregUtilTest, GetZerosVector) {
+  VectorType vty = VectorType::get({2, 4}, Builder().getI32Type());
+  TypedValue<VectorType> vec = getZerosVector(Builder(), vty);
+
+  EXPECT_THAT(vec, IsConstantOpWithSplatValue(vty, int32_t{0}));
+}
+
+TEST_F(VregUtilTest, GetZerosLikeVector) {
+  VectorType vty = VectorType::get({2, 4}, Builder().getF32Type());
+  TypedValue<VectorType> in_vec = Builder().create<vector::SplatOp>(
+      vty, Builder().create<arith::ConstantOp>(
+               vty.getElementType(), Builder().getF32FloatAttr(1.0f)));
+  TypedValue<VectorType> vec = getZerosLikeVector(Builder(), in_vec);
+
+  EXPECT_THAT(vec, IsConstantOpWithSplatValue(vty, float{0.0f}));
+}
+
+TEST_F(VregUtilTest, GetX32VmaskByPaddingEndDim0) {
+  constexpr std::array<int64_t, 2> kTargetShape = {4, 8};
+  FailureOr<TypedValue<VectorType>> vec = getX32VmaskByPaddingEnd(
+      Builder(), /*padding=*/1, /*target_shape=*/kTargetShape,
+      /*dim=*/0);
+  ASSERT_TRUE(succeeded(vec));
+
+  auto cmp_op = dyn_cast<arith::CmpIOp>(vec.value().getDefiningOp());
+  ASSERT_TRUE(cmp_op != nullptr);
+  EXPECT_EQ(cmp_op.getPredicate(), arith::CmpIPredicate::slt);
+
+  auto iota_op = dyn_cast<tpu::IotaOp>(cmp_op.getLhs().getDefiningOp());
+  ASSERT_TRUE(iota_op != nullptr);
+  EXPECT_THAT(iota_op.getDimension(), Optional(Eq(0)));
+
+  EXPECT_THAT(
+      cmp_op.getRhs(),
+      IsConstantOpWithSplatValue(
+          VectorType::get(kTargetShape, Builder().getI32Type()), int32_t{3}));
+}
+
+TEST_F(VregUtilTest, GetX32VmaskByPaddingEndDim1) {
+  constexpr std::array<int64_t, 2> kTargetShape = {4, 8};
+  FailureOr<TypedValue<VectorType>> vec = getX32VmaskByPaddingEnd(
+      Builder(), /*padding=*/3, /*target_shape=*/kTargetShape,
+      /*dim=*/1);
+  ASSERT_TRUE(succeeded(vec));
+
+  auto cmp_op = dyn_cast<arith::CmpIOp>(vec.value().getDefiningOp());
+  ASSERT_TRUE(cmp_op != nullptr);
+  EXPECT_EQ(cmp_op.getPredicate(), arith::CmpIPredicate::slt);
+
+  auto iota_op = dyn_cast<tpu::IotaOp>(cmp_op.getLhs().getDefiningOp());
+  ASSERT_TRUE(iota_op != nullptr);
+  EXPECT_THAT(iota_op.getDimension(), Optional(Eq(1)));
+
+  EXPECT_THAT(
+      cmp_op.getRhs(),
+      IsConstantOpWithSplatValue(
+          VectorType::get(kTargetShape, Builder().getI32Type()), int32_t{5}));
+}
+
+}  // namespace
+
+}  // namespace mlir::tpu
diff --git a/jaxlib/mosaic/python/mosaic_gpu.py b/jaxlib/mosaic/python/mosaic_gpu.py
index f99f53cfdb69..cce2909be206 100644
--- a/jaxlib/mosaic/python/mosaic_gpu.py
+++ b/jaxlib/mosaic/python/mosaic_gpu.py
@@ -33,4 +33,5 @@
   from mlir.dialects._ods_common import _cext  # type: ignore[import-not-found]
 
 
-_cext.globals.append_dialect_search_prefix("jax.jaxlib.mosaic.python")
+# Add the parent module to the search prefix
+_cext.globals.append_dialect_search_prefix(__name__[:__name__.rfind(".")])
diff --git a/jaxlib/rocm/BUILD b/jaxlib/rocm/BUILD
index 1076f9a77bf8..6b481682a885 100644
--- a/jaxlib/rocm/BUILD
+++ b/jaxlib/rocm/BUILD
@@ -135,6 +135,44 @@ pybind_extension(
     ],
 )
 
+cc_library(
+    name = "miopen_rnn_kernels",
+    srcs = ["//jaxlib/gpu:rnn_kernels.cc"],
+    hdrs = ["//jaxlib/gpu:rnn_kernels.h"],
+    deps = [
+        ":hip_gpu_kernel_helpers",
+        ":hip_vendor",
+        "//jaxlib:handle_pool",
+        "//jaxlib:kernel_helpers",
+        "@xla//xla/service:custom_call_status",
+        "@local_config_rocm//rocm:miopen",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+        "@local_config_rocm//rocm:rocm_headers",
+    ],
+)
+
+pybind_extension(
+    name = "_rnn",
+    srcs = ["//jaxlib/gpu:rnn.cc"],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+    ],
+    features = ["-use_header_modules"],
+    module_name = "_rnn",
+    deps = [
+        ":hip_vendor",
+        ":miopen_rnn_kernels",
+        "//jaxlib:absl_status_casters",
+        "//jaxlib:kernel_nanobind_helpers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings:str_format",
+        "@nanobind",
+    ],
+)
+
 cc_library(
     name = "hip_solver_handle_pool",
     srcs = ["//jaxlib/gpu:solver_handle_pool.cc"],
@@ -501,6 +539,7 @@ py_library(
         ":_hybrid",
         ":_linalg",
         ":_prng",
+        ":_rnn",
         ":_solver",
         ":_sparse",
         ":_triton",
diff --git a/jaxlib/setup.py b/jaxlib/setup.py
index 4370aa3176aa..c2efd3d7b7a7 100644
--- a/jaxlib/setup.py
+++ b/jaxlib/setup.py
@@ -83,6 +83,7 @@ def has_ext_modules(self):
             'cuda/*',
             'cuda/nvvm/libdevice/libdevice*',
             'mosaic/*.py',
+            'mosaic/dialect/gpu/*.py',
             'mosaic/gpu/*.so',
             'mosaic/python/*.py',
             'mosaic/python/*.so',
diff --git a/jaxlib/tools/build_gpu_kernels_wheel.py b/jaxlib/tools/build_gpu_kernels_wheel.py
index 36c1b4d2cbfc..65412f0365dc 100644
--- a/jaxlib/tools/build_gpu_kernels_wheel.py
+++ b/jaxlib/tools/build_gpu_kernels_wheel.py
@@ -140,12 +140,13 @@ def prepare_wheel_rocm(
   copy_runfiles(
       dst_dir=plugin_dir,
       src_files=[
-          f"__main__/jaxlib/rocm/_solver.{pyext}",
           f"__main__/jaxlib/rocm/_blas.{pyext}",
           f"__main__/jaxlib/rocm/_linalg.{pyext}",
           f"__main__/jaxlib/rocm/_prng.{pyext}",
+          f"__main__/jaxlib/rocm/_solver.{pyext}",
           f"__main__/jaxlib/rocm/_sparse.{pyext}",
           f"__main__/jaxlib/rocm/_hybrid.{pyext}",
+          f"__main__/jaxlib/rocm/_rnn.{pyext}",
           f"__main__/jaxlib/rocm/_triton.{pyext}",
           f"__main__/jaxlib/rocm_plugin_extension.{pyext}",
           "__main__/jaxlib/version.py",
diff --git a/jaxlib/tools/build_wheel.py b/jaxlib/tools/build_wheel.py
index b46a50961169..4b71bd5de2d8 100644
--- a/jaxlib/tools/build_wheel.py
+++ b/jaxlib/tools/build_wheel.py
@@ -218,6 +218,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
       dst_dir=mosaic_python_dir,
       src_files=[
           "__main__/jaxlib/mosaic/python/layout_defs.py",
+          "__main__/jaxlib/mosaic/python/mosaic_gpu.py",
           "__main__/jaxlib/mosaic/python/tpu.py",
       ],
   )
@@ -225,6 +226,16 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
   patch_copy_mlir_import(
       "__main__/jaxlib/mosaic/python/_tpu_gen.py", dst_dir=mosaic_python_dir
   )
+  mosaic_gpu_dir = jaxlib_dir / "mosaic" / "dialect" / "gpu"
+  os.makedirs(mosaic_gpu_dir)
+  patch_copy_mlir_import(
+      "__main__/jaxlib/mosaic/dialect/gpu/_mosaic_gpu_gen_ops.py",
+      dst_dir=mosaic_gpu_dir,
+  )
+  patch_copy_mlir_import(
+      "__main__/jaxlib/mosaic/dialect/gpu/_mosaic_gpu_gen_enums.py",
+      dst_dir=mosaic_gpu_dir,
+  )
 
   copy_runfiles(
       dst_dir=jaxlib_dir / "mlir",
@@ -316,6 +327,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
           f"__main__/jaxlib/mlir/_mlir_libs/_mlirHlo.{pyext}",
           f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsSparseTensor.{pyext}",
           f"__main__/jaxlib/mlir/_mlir_libs/_mlirSparseTensorPasses.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mosaic_gpu_ext.{pyext}",
           f"__main__/jaxlib/mlir/_mlir_libs/_tpu_ext.{pyext}",
           f"__main__/jaxlib/mlir/_mlir_libs/_sdy.{pyext}",
           f"__main__/jaxlib/mlir/_mlir_libs/_stablehlo.{pyext}",
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
index 5532fdf0303f..edf07f1bd5bf 100644
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -219,6 +219,29 @@ def f(x):
          [ 1  2  3  4  5  6  7  8  9 10 12 13 14]]
     """))
 
+  def test_debug_print_respects_numpy_printoptions(self):
+    def f(x):
+      with np.printoptions(precision=2, suppress=True):
+        jax.debug.print("{}", x)
+    x = np.array([1.2345, 2.3456, 1E-7])
+
+    # Default numpy print options:
+    with jtu.capture_stdout() as output:
+      jax.debug.print("{}", x)
+    self.assertEqual(output(), "[1.2345e+00 2.3456e+00 1.0000e-07]\n")
+
+    # Modified print options without JIT:
+    with jtu.capture_stdout() as output:
+      f(x)
+      jax.effects_barrier()
+    self.assertEqual(output(), "[1.23 2.35 0.  ]\n")
+
+    # Modified print options with JIT:
+    with jtu.capture_stdout() as output:
+      jax.jit(f)(x)
+      jax.effects_barrier()
+    self.assertEqual(output(), "[1.23 2.35 0.  ]\n")
+
 
 class DebugPrintTransformationTest(jtu.JaxTestCase):
 
diff --git a/tests/experimental_rnn_test.py b/tests/experimental_rnn_test.py
index d886a84f914b..524920fe9ec9 100644
--- a/tests/experimental_rnn_test.py
+++ b/tests/experimental_rnn_test.py
@@ -34,15 +34,19 @@ class RnnTest(jtu.JaxTestCase):
       num_layers=[1, 4],
       bidirectional=[True, False],
   )
-  @jtu.run_on_devices("cuda")
+  @jtu.run_on_devices("cuda", "rocm")
   @jax.default_matmul_precision("float32")
   def test_lstm(self, batch_size: int, seq_len: int, input_size: int,
                 hidden_size: int, num_layers: int, bidirectional: bool):
     # TODO(phawkins): Partially disable this on cudnn version per b/281071013
-    if (batch_size == 1 and seq_len == 4 and input_size == 1 and
+    if jtu.is_device_cuda() and (batch_size == 1 and seq_len == 4 and input_size == 1 and
         hidden_size == 6 and num_layers == 4 and bidirectional == False):
       self.skipTest("Test requires cudnn >= 8.8")
 
+    # TODO(ruturaj4): Bidirectional doesn't quite work well with rocm.
+    if bidirectional and jtu.is_device_rocm():
+      self.skipTest("Bidirectional mode is not available for ROCm.")
+
     num_directions = 2 if bidirectional else 1
     seq_length_key, root_key = jax.random.split(jax.random.PRNGKey(0))
 
@@ -61,6 +65,7 @@ def test_lstm(self, batch_size: int, seq_len: int, input_size: int,
     weights = rnn.init_lstm_weight(k4, input_size, hidden_size, num_layers,
                                    bidirectional)
     def f(weights, x, h_0, c_0):
+      weights = rnn.swap_lstm_gates(weights, input_size, hidden_size, num_layers, bidirectional)
       y, h, c = rnn.lstm(
         x,
         h_0,
diff --git a/tests/fused_attention_stablehlo_test.py b/tests/fused_attention_stablehlo_test.py
index 95ec4ce72eb4..c5cfb9d7daf7 100644
--- a/tests/fused_attention_stablehlo_test.py
+++ b/tests/fused_attention_stablehlo_test.py
@@ -254,8 +254,6 @@ def dot_product_attention_fp8(query, key, value, fp8_metas):
 class DotProductAttentionTest(jtu.JaxTestCase):
   def setUp(self):
     super().setUp()
-    if jax.device_count() < 4:
-      self.skipTest("Requires more than 4 devices.")
     try:
       cudnn_version = check_cudnn_version()
     except RuntimeError as e:
@@ -366,6 +364,8 @@ def test_sdpa(self, batch_size: int, seq_len: int, num_heads: int,
 
   @jtu.run_on_devices("cuda")
   def test_sdpa_inference(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
     query = jax.random.normal(
         k1, (4, 1024, 4, 64), dtype=jnp.bfloat16)
@@ -407,6 +407,8 @@ def test_sdpa_inference(self):
 
   @jtu.run_on_devices("cuda")
   def test_sdpa_var_seq(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     self.skipTest("Skip before fixed.")
     k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4)
     query = jax.random.normal(
@@ -438,6 +440,8 @@ def test_sdpa_var_seq(self):
 
   @jtu.run_on_devices("cuda")
   def test_sdpa_broadcast_bias_and_dbias(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     try:
       cudnn_version = check_cudnn_version()
     except RuntimeError as e:
@@ -504,6 +508,8 @@ def test_sdpa_broadcast_bias_and_dbias(self):
   )
   @jtu.run_on_devices("cuda")
   def test_sdpa_dbias(self, batch_size: int):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     # cuDNN only supports dbias when batch size is 1. If the batch size is
     # greater, dbias is silently set to all zeros. This test verifies this
     # behavior for both vmap and regular use cases.
@@ -540,6 +546,8 @@ def attn_vjp(x, bias, mask, target_fn):
 
   @jtu.run_on_devices("cuda")
   def test_sdpa_sliding_window_length(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4)
     query = jax.random.normal(
         k1, (4, 1024, 4, 64), dtype=jnp.bfloat16)
@@ -571,8 +579,43 @@ def test_sdpa_sliding_window_length(self):
     self.assertArraysAllClose(key_grad_ref, key_grad, rtol=1e-5, atol=1e-5)
     self.assertArraysAllClose(value_grad_ref, value_grad, rtol=1e-5, atol=1e-5)
 
+  @jtu.run_on_devices("cuda")
+  def test_sdpa_large_head_size(self):
+    try:
+      cudnn_version = check_cudnn_version()
+    except RuntimeError as e:
+      self.skipTest(str(e))
+      return
+    if cudnn_version < 90500:
+      self.skipTest("Requires >= cuDNN 9.5.0")
+    if not jtu.is_cuda_compute_capability_at_least("9.0"):
+      self.skipTest("Requires at least Hopper arch")
+
+    B, T, N, H = 2, 64, 2, 256
+    bf16 = jnp.bfloat16
+    keys = jax.random.split(jax.random.key(0), 4)
+    query = jax.random.normal(keys[0], (B, T, N, H), dtype=bf16)
+    key = jax.random.normal(keys[1], (B, T, N, H), dtype=bf16)
+    value = jax.random.normal(keys[2], (B, T, N, H), dtype=bf16)
+    grad = jax.random.normal(keys[3], (B, T, N, H), dtype=bf16)
+    sdpa_train_ans = jax.jit(partial(
+        sdpa_train, scale=1.0, mask_type=MaskType.CAUSAL, dropout_rate=0)
+    )
+    sdpa_train_rfc = jax.jit(partial(
+        sdpa_train_ref, scale=1.0, mask_type=MaskType.CAUSAL, dropout_rate=0)
+    )
+
+    out_ans, grads_ans = sdpa_train_ans(query, key, value, grad, None, None)
+    out_ref, grads_ref = sdpa_train_rfc(query, key, value, grad, None, None)
+    self.assertArraysAllClose(out_ref, out_ans)
+    self.assertArraysAllClose(grads_ref[0], grads_ans[0])
+    self.assertArraysAllClose(grads_ref[1], grads_ans[1])
+    self.assertArraysAllClose(grads_ref[2], grads_ans[2])
+
   @jtu.run_on_devices("cuda")
   def test_layouts(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     dtype = "bfloat16"
     B, T, N, H = 4, 1024, 8, 128
     S = T
@@ -600,6 +643,8 @@ def _cvt_back(x):
     self.assertArraysAllClose(dv_ref, _cvt_back(dv))
 
   def test_sdpa_utils(self):
+    if jax.device_count() < 4:
+      self.skipTest("Requires more than 4 devices.")
     test_cases = [
       (1, 257, 64, 8905, False, True, True),
       (1, 1024, 64, 8905, False, False, True),
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 33e8fd15f9f7..9db2f5bccca4 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -4386,7 +4386,7 @@ def regions_with_inaccuracies_keep(*to_keep):
       regions_with_inaccuracies_keep('q1.real', 'q2.real', 'q3.real', 'q4.real', 'ninf.imag', 'pinf.imag', 'ninfj.imag', 'pinfj.imag')
 
     elif name == 'log10':
-      regions_with_inaccuracies_keep('q1', 'q2', 'q3', 'q4', 'ninf.imag', 'pinf.imag', 'ninfj.imag', 'pinfj.imag', 'zero.imag')
+      regions_with_inaccuracies_keep('q1.real', 'q2.real', 'q3.real', 'q4.real', 'ninf.imag', 'pinf.imag', 'ninfj.imag', 'pinfj.imag')
 
     elif name == 'exp':
       regions_with_inaccuracies_keep('pos.imag', 'pinf.imag', 'mpos.imag')
diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
index d501f926736c..594df44e1db6 100644
--- a/tests/pallas/BUILD
+++ b/tests/pallas/BUILD
@@ -39,8 +39,8 @@ jax_multiplatform_test(
         "tpu",
     ],
     enable_configs = [
-        "gpu_a100_x32",
-        "gpu_h100_x32",
+        "gpu_a100",
+        "gpu_h100",
     ],
     shard_count = {
         "cpu": 8,
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
index 6e4928082ac6..373388d97691 100644
--- a/tests/pallas/pallas_test.py
+++ b/tests/pallas/pallas_test.py
@@ -1688,8 +1688,8 @@ def cond(state):
 
       def body(state):
         i, s = state
-        sl = jax.lax.div(i, 128)
-        l = jax.lax.rem(i, 128)
+        sl = jax.lax.div(i, jnp.astype(128, i.dtype))
+        l = jax.lax.rem(i, jnp.astype(128, i.dtype))
         v = pl.load(x_ref, (0, sl, l))
         return i + 1, s + v
 
diff --git a/tests/pallas/tpu_pallas_test.py b/tests/pallas/tpu_pallas_test.py
index 9c4788d7447f..8eee883a536e 100644
--- a/tests/pallas/tpu_pallas_test.py
+++ b/tests/pallas/tpu_pallas_test.py
@@ -2555,9 +2555,7 @@ def kernel(x_ref, out_ref):
 
     np.testing.assert_array_equal(out, np.reshape(x, (8, 128)))
 
-  @only_passes_in_interpret()
   def test_retiling2(self):
-    """b/348040767"""
     x = np.arange(1 * 8 * 1024, dtype=jnp.bfloat16).reshape(1, 8, 1024)
 
     def kernel(x_ref, out_ref):
diff --git a/tests/pgle_test.py b/tests/pgle_test.py
index dbf67b1421ab..05154ae0376c 100644
--- a/tests/pgle_test.py
+++ b/tests/pgle_test.py
@@ -95,8 +95,8 @@ def testPGLEProfilerGetFDOProfileLarge(self):
         out_shardings=NamedSharding(mesh, PartitionSpec('x')),
         compiler_options={
             'xla_gpu_enable_latency_hiding_scheduler': 'True',
-            # TODO(patrios): Remove this flag once b/376647494 is fixed.
-            'xla_gpu_graph_min_graph_size': '100000',
+            # TODO(b/37664749): Remove this flag once the bug is fixed.
+            'xla_gpu_enable_command_buffer': '',
         },
     )
     def f(x):
@@ -133,8 +133,6 @@ def testAutoPgle(self):
           out_shardings=NamedSharding(mesh, PartitionSpec('x')),
           compiler_options={
               'xla_gpu_enable_latency_hiding_scheduler': 'True',
-              # TODO(patrios): Remove this flag once b/376647494 is fixed.
-              'xla_gpu_graph_min_graph_size': '100000',
               'xla_dump_to': dump_dir,
               'xla_gpu_experimental_dump_fdo_profiles': 'True'
           },
@@ -217,8 +215,6 @@ def testAutoPgleWithPersistentCache(self):
           out_shardings=NamedSharding(mesh, PartitionSpec('x')),
           compiler_options={
               'xla_gpu_enable_latency_hiding_scheduler': 'True',
-              # TODO(patrios): Remove this flag once b/376647494 is fixed.
-              'xla_gpu_graph_min_graph_size': '100000',
               'xla_dump_to': dump_dir,
               'xla_gpu_experimental_dump_fdo_profiles': 'True'
           },
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index a24f75102825..e4f62a324d54 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -2207,6 +2207,23 @@ def f(x):
   #
   #   f(x)  # don't crash
 
+  def test_partial_auto_of_random_keys(self):
+    if config.use_shardy_partitioner.value:
+      self.skipTest('Shardy does not support full-to-shard.')
+
+    mesh = jtu.create_mesh((4, 2), ('i', 'j'))
+    keys = jax.random.split(jax.random.key(0), 8)
+
+    @jax.jit
+    def f(x):
+      return shard_map(lambda k: k,
+                       mesh, in_specs=P('i'), out_specs=P('i'),
+                       check_rep=False, auto=frozenset({'j'}))(keys)
+
+    y = f(keys) # don't crash
+    self.assertAllClose(jax.random.key_data(y), jax.random.key_data(keys),
+                        check_dtypes=False)
+
   def test_vmap_grad_shmap_spmd_axis_name_residuals(self):
     # https://github.com/jax-ml/jax/pull/21032
     mesh = jtu.create_mesh((4, 2), ('i', 'j'))
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index c88e78f1c053..1960c3a9d06f 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 = "dc7aaf834a0bb5a543f6cf98626284783a4a921c"
-XLA_SHA256 = "eda76cce64b33c00139120d6b4d4c2167d9f99dc957da54225a67ddb7ec7cb23"
+XLA_COMMIT = "ac6e71fe0cf864eec152de5ba761b76d8bef3153"
+XLA_SHA256 = "2b568ff365bc4b5c2b257002aa71f094a2b60357ceb1f2a1c6c33f4ad1a411bd"
 
 def repo():
     tf_http_archive(