Repository: google/jax
Pull request #22519 (2024-07-18 19:56:14)

Support `pathlib.Path` input to `start_trace`
If you want tests for this or anything else, let me know. 

sharadmv: A test would be nice, maybe add one [here](https://github.com/google/jax/blob/3dd0d74f2b54c89e5f919cd00c58683c08c89f08/tests/profiler_test.py)?

Patch:
diff --git a/docs/jax.extend.ffi.rst b/docs/jax.extend.ffi.rst
index ac8e38c5e..5928189eb 100644
--- a/docs/jax.extend.ffi.rst
+++ b/docs/jax.extend.ffi.rst
@@ -9,4 +9,3 @@
   ffi_call
   ffi_lowering
   pycapsule
-  register_ffi_target
diff --git a/docs/pallas/tpu/matmul.ipynb b/docs/pallas/tpu/matmul.ipynb
index b0af0f2a8..ec74ca274 100644
--- a/docs/pallas/tpu/matmul.ipynb
+++ b/docs/pallas/tpu/matmul.ipynb
@@ -6,7 +6,7 @@
     "id": "-z6pOJwvn-_j"
    },
    "source": [
-    "# Matrix Multiplication\n",
+    "# Pallas TPU - Your First Matmul\n",
     "\n",
     "In this guide, we'll write a matrix multiplication routine using Pallas. We'll also go over how to think about matmul performance on TPU and how to template a matmul kernel to fuse in operations."
    ]
@@ -232,9 +232,6 @@
     "\n",
     "Let's think about how to analyze matrix multiplication performance.\n",
     "\n",
-    "> Quick note: in this guide, we'll be discussing floating point operations, but want to make the distinction between FLOPs vs FLOP/s.\n",
-    "  When we say \"FLOPs\" we mean \"floating point operations\", as in a number of operations. When we say \"FLOP/s\", we refer to \"floating point operations *per second*\", as in a *rate* of performing floating point operations.\n",
-    "\n",
     "The number of FLOPs in a `(m, k) x (k, n)` matrix multiplication are (approximately) `2 * m * k * n`. (Technically it is `n * m * (2k - 1)` but for large enough `k` our approximation is sufficient.)\n",
     "\n",
     "The minimum amount of memory usage for a matrix multiply (assuming float32) is the total size of the inputs (copying into VMEM) plus the size of the output (copying into HBM). Thus the minimum bandwidth usage is `(m * k + k * n + m * n) * 4 bytes/float32`. Memory usage can be greater if we re-read the inputs multiple times, which is often the case.\n",
@@ -276,9 +273,9 @@
     "id": "agCtb2GMQazl"
    },
    "source": [
-    "Now that we can calculate the total number of FLOPs and (minimum) memory bandwidth usage of a matrix multiplication, let's now see how many FLOP/s we can execute and how much memory bandwidth we actually have on a particular TPU.\n",
+    "Now that we can calculate the amount of FLOPs and (minimum) memory bandwidth usage of a matrix multiplication, let's now see how many FLOP/s we can execute and how much memory bandwidth we actually have on a particular TPU.\n",
     "\n",
-    "This notebook was run on TPU v5e so we will utilize numbers for that specific chip (if you are running this notebook, your numbers may differ). TPU v5es have [197 TFLOP/s of bf16 compute and 819 GB/s of memory bandwidth](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v5e). We can compute the minimum arithmetic intensity (ratio of FLOP/s to memory bandwidth) needed to be compute bound in this analytical model on this chip by taking their ratio (about 240 FLOPs/byte on TPU v5e)."
+    "This notebook was run on TPU v5e so we will utilize numbers for that specific chip (if you are running this notebook, your numbers may differ). TPU v5es have [197 TFLOPs of bf16 compute and 819 GB/s of memory bandwidth](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v5e). We can compute the minimum arithmetic intensity (ratio of FLOPs to memory bandwidth) needed to be compute bound in this analytical model on this chip by taking their ratio (about 240 FLOPs/byte on TPU v5e)."
    ]
   },
   {
@@ -454,8 +451,6 @@
     "        grid=(m // bm, n // bn, k // bk),\n",
     "      ),\n",
     "      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n",
-    "      compiler_params=dict(mosaic=dict(\n",
-    "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\"))),\n",
     "  )(x, y)"
    ]
   },
@@ -482,9 +477,9 @@
    "source": [
     "## Performance of pipelined kernels\n",
     "\n",
-    "Our above analysis about FLOPs vs memory usage applies at a coarse scale i.e. when we are looking at the the size of a the total matrix multiplication. However, remember that in practice, we are pipelining the execution of a blocked matrix multiplication, meaning we have a loop in which we are doing matrix multiplication with smaller blocks.\n",
+    "Our above analysis about FLOPS vs memory usage applies at a coarse scale i.e. when we are looking at the the size of a the total matrix multiplication. However, remember that in practice, we are pipelining the execution of a blocked matrix multiplication, meaning we have a loop in which we are doing matrix multiplication with smaller blocks.\n",
     "\n",
-    "This means that we actually care about the FLOPs vs memory bandwidth usage of each individual instance of the kernel, not the global FLOPs vs memory bandwidth usage. Therefore, the block sizes `bm`, `bk`, `bn` are extremely important for performance. Even if we have the largest matrices in the world, if we pick very small `bm`, `bk`, and `bn`, we will be memory bound because each time we invoke the kernel we will have too few FLOPs to hide the memory transfers happening in the background.\n",
+    "This means that we actually care about the FLOPS vs memory bandwidth usage of each individual instance of the kernel, not the global FLOPS vs memory bandwidth usage. Therefore, the block sizes `bm`, `bk`, `bn` are extremely important for performance. Even if we have the largest matrices in the world, if we pick very small `bm`, `bk`, and `bn`, we will be memory bound because each time we invoke the kernel we will have too few FLOPs to hide the memory transfers happening in the background.\n",
     "\n",
     "The intuition should therefore be: to be compute bound, make the blocks as big as possible! There are two main constraints:\n",
     "\n",
@@ -512,33 +507,33 @@
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00029766598949208854\n",
       "Matmul FLOP/s:  7214407167121.377\n",
-      "FLOP/s utilization: 3.6621%\n",
+      "FLOPs utilization: 3.6621%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.011771515250438824\n",
       "Matmul FLOP/s:  11675553278230.387\n",
-      "FLOP/s utilization: 5.9267%\n",
+      "FLOPs utilization: 5.9267%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.09183577066054567\n",
       "Matmul FLOP/s:  11972585626140.668\n",
-      "FLOP/s utilization: 6.0775%\n",
+      "FLOPs utilization: 6.0775%\n",
       "\n",
       "================bm=512, bk=1024, bn=1024===================\n",
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00012708659982308746\n",
       "Matmul FLOP/s:  16897797651282.135\n",
-      "FLOP/s utilization: 8.5776%\n",
+      "FLOPs utilization: 8.5776%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.00088908776990138\n",
       "Matmul FLOP/s:  154584235803001.88\n",
-      "FLOP/s utilization: 78.4692%\n",
+      "FLOPs utilization: 78.4692%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.006099433819763363\n",
       "Matmul FLOP/s:  180264539343531.62\n",
-      "FLOP/s utilization: 91.5048%\n",
+      "FLOPs utilization: 91.5048%\n",
       "\n"
      ]
     }
@@ -567,7 +562,7 @@
     "  print(\"Matmul time: \", time)\n",
     "  mm_flops = matmul_flops(m, k, n) / time\n",
     "  print(\"Matmul FLOP/s: \", mm_flops)\n",
-    "  print(f\"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
+    "  print(f\"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
     "  print()\n",
     "\n",
     "print(\"================bm=128, bk=128, bn=128===================\")\n",
@@ -611,17 +606,17 @@
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00011943008983507753\n",
       "Matmul FLOP/s:  17981093801113.996\n",
-      "FLOP/s utilization: 9.1275%\n",
+      "FLOPs utilization: 9.1275%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.0008272899803705514\n",
       "Matmul FLOP/s:  166131533963991.34\n",
-      "FLOP/s utilization: 84.3307%\n",
+      "FLOPs utilization: 84.3307%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.006047147869830951\n",
       "Matmul FLOP/s:  181823175395037.44\n",
-      "FLOP/s utilization: 92.2960%\n",
+      "FLOPs utilization: 92.2960%\n",
       "\n"
      ]
     }
@@ -727,8 +722,6 @@
     "        grid=(m // bm, n // bn, k // bk),\n",
     "      ),\n",
     "      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n",
-    "      compiler_params=dict(mosaic=dict(\n",
-    "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\"))),\n",
     "  )(x, y)"
    ]
   },
@@ -762,33 +755,33 @@
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.0003029372810851783\n",
       "Matmul FLOP/s:  7088872126624.065\n",
-      "FLOP/s utilization: 3.5984%\n",
+      "FLOPs utilization: 3.5984%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.012017967159627005\n",
       "Matmul FLOP/s:  11436123235026.848\n",
-      "FLOP/s utilization: 5.8051%\n",
+      "FLOPs utilization: 5.8051%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.09500920018996112\n",
       "Matmul FLOP/s:  11572685861765.383\n",
-      "FLOP/s utilization: 5.8745%\n",
+      "FLOPs utilization: 5.8745%\n",
       "\n",
       "================bm=512, bk=1024, bn=1024===================\n",
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00012131539988331496\n",
       "Matmul FLOP/s:  17701657415839.363\n",
-      "FLOP/s utilization: 8.9856%\n",
+      "FLOPs utilization: 8.9856%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.0008790623804088682\n",
       "Matmul FLOP/s:  156347213275211.03\n",
-      "FLOP/s utilization: 79.3641%\n",
+      "FLOPs utilization: 79.3641%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.006107717020204291\n",
       "Matmul FLOP/s:  180020067095253.78\n",
-      "FLOP/s utilization: 91.3807%\n",
+      "FLOPs utilization: 91.3807%\n",
       "\n"
      ]
     }
@@ -811,7 +804,7 @@
     "  print(\"Matmul time: \", time)\n",
     "  mm_flops = matmul_flops(m, k, n) / time\n",
     "  print(\"Matmul FLOP/s: \", mm_flops)\n",
-    "  print(f\"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
+    "  print(f\"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
     "  print()\n",
     "\n",
     "print(\"================bm=128, bk=128, bn=128===================\")\n",
@@ -936,33 +929,33 @@
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00030103540048003196\n",
       "Matmul FLOP/s:  7133658182976.541\n",
-      "FLOP/s utilization: 3.6211%\n",
+      "FLOPs utilization: 3.6211%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.011807117109419778\n",
       "Matmul FLOP/s:  11640348122095.826\n",
-      "FLOP/s utilization: 5.9088%\n",
+      "FLOPs utilization: 5.9088%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.09181861146935262\n",
       "Matmul FLOP/s:  11974823079773.941\n",
-      "FLOP/s utilization: 6.0786%\n",
+      "FLOPs utilization: 6.0786%\n",
       "\n",
       "================bm=512, bk=1024, bn=1024===================\n",
       "----- 1024 x 1024 x 1024 -----\n",
       "Matmul time:  0.00012622540001757442\n",
       "Matmul FLOP/s:  17013086492108.6\n",
-      "FLOP/s utilization: 8.6361%\n",
+      "FLOPs utilization: 8.6361%\n",
       "\n",
       "----- 4096 x 4096 x 4096 -----\n",
       "Matmul time:  0.000896632740041241\n",
       "Matmul FLOP/s:  153283442968721.44\n",
-      "FLOP/s utilization: 77.8089%\n",
+      "FLOPs utilization: 77.8089%\n",
       "\n",
       "----- 8192 x 8192 x 8192 -----\n",
       "Matmul time:  0.006130605939542875\n",
       "Matmul FLOP/s:  179347953304919.88\n",
-      "FLOP/s utilization: 91.0396%\n",
+      "FLOPs utilization: 91.0396%\n",
       "\n"
      ]
     }
@@ -986,7 +979,7 @@
     "  print(\"Matmul time: \", time)\n",
     "  mm_flops = matmul_flops(m, k, n) / time\n",
     "  print(\"Matmul FLOP/s: \", mm_flops)\n",
-    "  print(f\"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
+    "  print(f\"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%\")\n",
     "  print()\n",
     "\n",
     "\n",
diff --git a/docs/pallas/tpu/matmul.md b/docs/pallas/tpu/matmul.md
index b01d8b82c..e4858f558 100644
--- a/docs/pallas/tpu/matmul.md
+++ b/docs/pallas/tpu/matmul.md
@@ -14,7 +14,7 @@ kernelspec:
 
 +++ {"id": "-z6pOJwvn-_j"}
 
-# Matrix Multiplication
+# Pallas TPU - Your First Matmul
 
 In this guide, we'll write a matrix multiplication routine using Pallas. We'll also go over how to think about matmul performance on TPU and how to template a matmul kernel to fuse in operations.
 
@@ -180,9 +180,6 @@ np.testing.assert_array_equal(x @ y, matmul(x, y))
 
 Let's think about how to analyze matrix multiplication performance.
 
-> Quick note: in this guide, we'll be discussing floating point operations, but want to make the distinction between FLOPs vs FLOP/s.
-  When we say "FLOPs" we mean "floating point operations", as in a number of operations. When we say "FLOP/s", we refer to "floating point operations *per second*", as in a *rate* of performing floating point operations.
-
 The number of FLOPs in a `(m, k) x (k, n)` matrix multiplication are (approximately) `2 * m * k * n`. (Technically it is `n * m * (2k - 1)` but for large enough `k` our approximation is sufficient.)
 
 The minimum amount of memory usage for a matrix multiply (assuming float32) is the total size of the inputs (copying into VMEM) plus the size of the output (copying into HBM). Thus the minimum bandwidth usage is `(m * k + k * n + m * n) * 4 bytes/float32`. Memory usage can be greater if we re-read the inputs multiple times, which is often the case.
@@ -205,9 +202,9 @@ print(matmul_membw(1024, 1024, 1024, jnp.float32))
 
 +++ {"id": "agCtb2GMQazl"}
 
-Now that we can calculate the total number of FLOPs and (minimum) memory bandwidth usage of a matrix multiplication, let's now see how many FLOP/s we can execute and how much memory bandwidth we actually have on a particular TPU.
+Now that we can calculate the amount of FLOPs and (minimum) memory bandwidth usage of a matrix multiplication, let's now see how many FLOP/s we can execute and how much memory bandwidth we actually have on a particular TPU.
 
-This notebook was run on TPU v5e so we will utilize numbers for that specific chip (if you are running this notebook, your numbers may differ). TPU v5es have [197 TFLOP/s of bf16 compute and 819 GB/s of memory bandwidth](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v5e). We can compute the minimum arithmetic intensity (ratio of FLOP/s to memory bandwidth) needed to be compute bound in this analytical model on this chip by taking their ratio (about 240 FLOPs/byte on TPU v5e).
+This notebook was run on TPU v5e so we will utilize numbers for that specific chip (if you are running this notebook, your numbers may differ). TPU v5es have [197 TFLOPs of bf16 compute and 819 GB/s of memory bandwidth](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v5e). We can compute the minimum arithmetic intensity (ratio of FLOPs to memory bandwidth) needed to be compute bound in this analytical model on this chip by taking their ratio (about 240 FLOPs/byte on TPU v5e).
 
 ```{code-cell} ipython3
 :id: WUydNX2-K6Oy
@@ -311,8 +308,6 @@ def matmul(
         grid=(m // bm, n // bn, k // bk),
       ),
       out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-      compiler_params=dict(mosaic=dict(
-          dimension_semantics=("parallel", "parallel", "arbitrary"))),
   )(x, y)
 ```
 
@@ -330,9 +325,9 @@ np.testing.assert_array_equal(x @ y, matmul(x, y))
 
 ## Performance of pipelined kernels
 
-Our above analysis about FLOPs vs memory usage applies at a coarse scale i.e. when we are looking at the the size of a the total matrix multiplication. However, remember that in practice, we are pipelining the execution of a blocked matrix multiplication, meaning we have a loop in which we are doing matrix multiplication with smaller blocks.
+Our above analysis about FLOPS vs memory usage applies at a coarse scale i.e. when we are looking at the the size of a the total matrix multiplication. However, remember that in practice, we are pipelining the execution of a blocked matrix multiplication, meaning we have a loop in which we are doing matrix multiplication with smaller blocks.
 
-This means that we actually care about the FLOPs vs memory bandwidth usage of each individual instance of the kernel, not the global FLOPs vs memory bandwidth usage. Therefore, the block sizes `bm`, `bk`, `bn` are extremely important for performance. Even if we have the largest matrices in the world, if we pick very small `bm`, `bk`, and `bn`, we will be memory bound because each time we invoke the kernel we will have too few FLOPs to hide the memory transfers happening in the background.
+This means that we actually care about the FLOPS vs memory bandwidth usage of each individual instance of the kernel, not the global FLOPS vs memory bandwidth usage. Therefore, the block sizes `bm`, `bk`, `bn` are extremely important for performance. Even if we have the largest matrices in the world, if we pick very small `bm`, `bk`, and `bn`, we will be memory bound because each time we invoke the kernel we will have too few FLOPs to hide the memory transfers happening in the background.
 
 The intuition should therefore be: to be compute bound, make the blocks as big as possible! There are two main constraints:
 
@@ -370,7 +365,7 @@ def analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,
   print("Matmul time: ", time)
   mm_flops = matmul_flops(m, k, n) / time
   print("Matmul FLOP/s: ", mm_flops)
-  print(f"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
+  print(f"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
   print()
 
 print("================bm=128, bk=128, bn=128===================")
@@ -477,8 +472,6 @@ def matmul(
         grid=(m // bm, n // bn, k // bk),
       ),
       out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-      compiler_params=dict(mosaic=dict(
-          dimension_semantics=("parallel", "parallel", "arbitrary"))),
   )(x, y)
 ```
 
@@ -512,7 +505,7 @@ def analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,
   print("Matmul time: ", time)
   mm_flops = matmul_flops(m, k, n) / time
   print("Matmul FLOP/s: ", mm_flops)
-  print(f"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
+  print(f"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
   print()
 
 print("================bm=128, bk=128, bn=128===================")
@@ -628,7 +621,7 @@ def analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,
   print("Matmul time: ", time)
   mm_flops = matmul_flops(m, k, n) / time
   print("Matmul FLOP/s: ", mm_flops)
-  print(f"FLOP/s utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
+  print(f"FLOPs utilization: {mm_flops / v5e_bf16_flops * 100:.4f}%")
   print()
 
 
diff --git a/jax/_src/extend/ffi.py b/jax/_src/extend/ffi.py
index df1c09eff..428187d57 100644
--- a/jax/_src/extend/ffi.py
+++ b/jax/_src/extend/ffi.py
@@ -29,7 +29,6 @@ from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
 from jax._src.lib import jaxlib
-from jax._src.lib import xla_client
 from jax._src.lib.mlir import ir
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray
 import numpy as np
@@ -37,31 +36,6 @@ import numpy as np
 map, unsafe_map = util.safe_map, map
 
 
-def register_ffi_target(
-    name: str,
-    fn: Any,
-    platform: str = "cpu",
-    api_version: int = 1,
-    **kwargs: Any,
-) -> None:
-  """Registers a foreign function target.
-
-  Args:
-    name: the name of the target.
-    fn: a ``PyCapsule`` object containing the function pointer, or a ``dict``
-      where the keys are FFI stage names (e.g. `"execute"`) and the values are
-      ``PyCapsule`` objects continaing a pointer to the handler for that stage.
-    platform: the target platform.
-    api_version: the XLA custom call API version to use. Supported versions are:
-      1 (default) for the typed FFI or 0 for the earlier "custom call" API.
-    kwargs: any extra keyword arguments are passed directly to
-      :func:`~jaxlib.xla_client.register_custom_call_target` for more advanced
-      use cases.
-  """
-  return xla_client.register_custom_call_target(name, fn, platform, api_version,
-                                                **kwargs)
-
-
 def pycapsule(funcptr):
   """Wrap a ctypes function pointer in a PyCapsule.
 
diff --git a/jax/_src/pallas/mosaic/pipeline.py b/jax/_src/pallas/mosaic/pipeline.py
index 08ec28d71..0d778a60c 100644
--- a/jax/_src/pallas/mosaic/pipeline.py
+++ b/jax/_src/pallas/mosaic/pipeline.py
@@ -549,9 +549,9 @@ class Scheduler:
   # outputs, and in-out accumulators.
 
   def initialize(self, buffered_ref, src_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule["prologue_copy_in"](self, buffered_ref, src_ref)
+    pred = self.first_step_ever
+    if schedule is not None:
+      pred = schedule['prologue_copy_in'](self, buffered_ref, src_ref)
 
     with jax.named_scope("ep_initialize"):
       @pl.when(self.first_step_ever)
@@ -563,44 +563,36 @@ class Scheduler:
         if buffered_ref.is_input:
           buffered_ref.copy_in(src_ref, self.indices)
 
-      # In the prologue this makes it so we wait on the prologue copy to finish.
-      # In other iterations this is the regular swap.
       buffered_ref.swap_slots()
 
   def wait_in(self, buffered_ref, src_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule["wait_in"](self, buffered_ref, src_ref)
+    pred = self.has_changed(buffered_ref) | self.first_step
+    if schedule is not None:
+      pred = schedule['wait_in'](self, buffered_ref, src_ref)
 
     @jax.named_scope("ep_wait_in")
     def _wait():
       if buffered_ref.is_input:
         buffered_ref.wait_in(src_ref, self.indices)
       if buffered_ref.is_accumulator:
-        # In most cases we won't be waiting when init_accumulators is True,
-        # so this is usually just setting what we just copied.
         buffered_ref.set_accumulator(self.init_accumulators)
     @jax.named_scope("ep_set_accum")
     def _no_wait():
       if buffered_ref.is_accumulator:
-
-        @pl.when(self.first_step | self.has_changed(buffered_ref))
+        @pl.when(self.first_step)
         def _set_accumulator():
-          # In most cases we will skip waiting when init_accumulators is True,
-          # so this is usually just setting the accumulator to 0.
           buffered_ref.set_accumulator(self.init_accumulators)
     lax.cond(pred, _wait, _no_wait)
 
   def copy_in(self, buffered_ref, src_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule['copy_in'](self, buffered_ref, src_ref)
+    pred = self.will_change(buffered_ref) & ~self.last_step_ever
+    if schedule is not None:
+      pred = schedule['copy_in'](self, buffered_ref, src_ref)
 
     @pl.when(pred)
     @jax.named_scope("ep_copy_in")
     def _send():
       if buffered_ref.is_input:
-        # We skip the last step because that's what prefetch is for.
         @pl.when(~self.last_step)
         def _copy_in():
           buffered_ref.copy_in(src_ref, self.next_indices)
@@ -609,23 +601,24 @@ class Scheduler:
 
   # convenience method for prefetch callbacks.
   def prefetch(self, buffered_ref, src_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule['prefetch'](self, buffered_ref, src_ref)
+    pred = ((self.will_change(buffered_ref) | self.last_step) &
+            ~self.last_step_ever)
+    if schedule is not None:
+      pred = schedule['prefetch'](self, buffered_ref, src_ref)
 
     @pl.when(pred)
     @jax.named_scope("ep_prefetch")
     def _send():
       if buffered_ref.is_input:
-        # Prefetch should only run on the last step.
         @pl.when(self.last_step)
         def _prefetch_in():
           buffered_ref.copy_in(src_ref, self.next_indices)
 
   def wait_out(self, buffered_ref, dst_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule['wait_out'](self, buffered_ref, dst_ref)
+    pred = ((self.has_changed(buffered_ref) | self.first_step) &
+            ~self.first_step_ever)
+    if schedule is not None:
+      pred = schedule['wait_out'](self, buffered_ref, dst_ref)
 
     @pl.when(pred)
     @jax.named_scope("ep_wait_out")
@@ -637,9 +630,9 @@ class Scheduler:
   #     cycle
 
   def copy_out(self, buffered_ref, dst_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule['copy_out'](self, buffered_ref, dst_ref)
+    pred = self.will_change(buffered_ref) | self.last_step
+    if schedule is not None:
+      pred = schedule['copy_out'](self, buffered_ref, dst_ref)
 
     @jax.named_scope("ep_copy_out")
     def _copy_out_and_accumulate():
@@ -650,20 +643,15 @@ class Scheduler:
     @jax.named_scope("ep_accum")
     def _just_accumulate():
       if buffered_ref.is_accumulator:
-        # We accumulate on the last step because we will set the accumulator
-        # on the next first step. We can optimize this away if it becomes
-        # a problem, but it is probably not worth the complexity to support
-        # chains of different pipelines that want to reuse the accumulator with
-        # slightly different schedules.
         @pl.when(self.last_step)
         def _accumulate():
           buffered_ref.accumulate()
     lax.cond(pred, _copy_out_and_accumulate, _just_accumulate)
 
   def finalize(self, buffered_ref, dst_ref, schedule=None):
-    if schedule is None:
-      schedule = _default_schedule
-    pred = schedule['epilogue_wait_out'](self, buffered_ref, dst_ref)
+    pred = self.last_step_ever
+    if schedule is not None:
+      pred = schedule['epilogue_wait_out'](self, buffered_ref, dst_ref)
 
     @pl.when(pred)
     @jax.named_scope("ep_finalize")
@@ -688,18 +676,12 @@ class Scheduler:
 # that the source and target HBM Refs change with each cycle.
 _default_schedule = dict(
     prologue_copy_in=lambda s, bref, _: s.first_step_ever,
-    # We assume that the source ref changed for prefetch.
     wait_in=lambda s, bref, _: s.has_changed(bref) | s.first_step,
     copy_in=lambda s, bref, _: s.will_change(bref) & ~s.last_step_ever,
-    # We assume that the source ref changed. E.g. because of a CM DMA.
     prefetch=lambda s, bref, _: (
-        (s.will_change(bref) | s.last_step) & ~s.last_step_ever
-    ),
-    # We assume that the target ref changed. E.g. because of a CM DMA.
+        (s.will_change(bref) | s.last_step) & ~s.last_step_ever),
     wait_out=lambda s, bref, _: (
-        (s.has_changed(bref) | s.first_step) & ~s.first_step_ever
-    ),
-    # We assume that the target ref is changing. E.g. because of a CM DMA.
+        (s.has_changed(bref) | s.first_step) & ~s.first_step_ever),
     copy_out=lambda s, bref, _: s.will_change(bref) | s.last_step,
     epilogue_wait_out=lambda s, bref, _: s.last_step_ever,
 )
@@ -711,40 +693,15 @@ _default_schedule = dict(
 # pipeline schedule.
 _fixed_schedule = dict(
     prologue_copy_in=lambda s, bref, _: s.first_step_ever,
-    # We don't assume that the source ref changed for prefetch.
     wait_in=lambda s, bref, _: s.has_changed(bref) | s.first_step_ever,
     copy_in=lambda s, bref, _: s.will_change(bref) & ~s.last_step_ever,
-    # We don't assume that the source ref changed.
     prefetch=lambda s, bref, _: s.will_change(bref) & ~s.last_step_ever,
-    # We don't assume that the target ref changed.
     wait_out=lambda s, bref, _: s.has_changed(bref) & ~s.first_step_ever,
-    # We don't assume that the target ref is changing.
     copy_out=lambda s, bref, _: s.will_change(bref) | s.last_step_ever,
     epilogue_wait_out=lambda s, bref, _: s.last_step_ever,
 )
 
 
-def skip_input_copies_when_init_accumulators(schedule) -> Any:
-  """Skip input copies in schedule when init_accumulators is True."""
-  new_schedule = {**schedule}
-  for k in ["prologue_copy_in", "wait_in", "copy_in"]:
-
-    def new_pred(original_pred_fn, *a):
-      pred = original_pred_fn(*a)
-      if a[1].is_accumulator:
-        pred &= ~a[0].init_accumulators
-      return pred
-
-    new_schedule[k] = functools.partial(
-        new_pred,
-        schedule[k],
-    )
-  return new_schedule
-
-
-_default_schedule = skip_input_copies_when_init_accumulators(_default_schedule)
-_fixed_schedule = skip_input_copies_when_init_accumulators(_fixed_schedule)
-
 def get_pipeline_schedule(schedule) -> Any:
   """Retrieve a named pipeline schedule or pass through fully specified one."""
   predefined_schedules = {
diff --git a/jax/_src/pallas/mosaic/verification.py b/jax/_src/pallas/mosaic/verification.py
index df186d463..d7463d9ed 100644
--- a/jax/_src/pallas/mosaic/verification.py
+++ b/jax/_src/pallas/mosaic/verification.py
@@ -164,11 +164,11 @@ class PrintCtx:
     result = io.StringIO()
     result.write(f"#define NDEVICE {num_devices}\n")
     result.write("#define NCORE 1\n")
-    result.write(f"byte _buf_readers[{self.num_buffers}*NDEVICE*NCORE] = 0;\n")
+    result.write(f"int  _buf_readers[{self.num_buffers}*NDEVICE*NCORE] = 0;\n")
     result.write(f"bool _buf_written[{self.num_buffers}*NDEVICE*NCORE] = 0;\n")
-    result.write(f"byte _sems[{self.num_semaphores}*NDEVICE*NCORE] = 0;\n")
+    result.write(f"int  _sems[{self.num_semaphores}*NDEVICE*NCORE] = 0;\n")
     if has_barrier_sems:
-      result.write("byte _barrier_sems[NDEVICE*NCORE] = 0;\n")
+      result.write("int _barrier_sems[NDEVICE*NCORE] = 0;\n")
     result.write(PREAMBLE)
     result.write("\n")
     parallel_threads = math.prod(parallel_iteration_bounds)
diff --git a/jax/_src/profiler.py b/jax/_src/profiler.py
index cad4826ba..1330a2148 100644
--- a/jax/_src/profiler.py
+++ b/jax/_src/profiler.py
@@ -129,7 +129,7 @@ def start_trace(log_dir, create_perfetto_link: bool = False,
     _profile_state.create_perfetto_link = create_perfetto_link
     _profile_state.create_perfetto_trace = (
         create_perfetto_trace or create_perfetto_link)
-    _profile_state.log_dir = log_dir
+    _profile_state.log_dir = str(log_dir)
 
 
 def _write_perfetto_trace_file(log_dir):
diff --git a/jax/_src/xla_bridge.py b/jax/_src/xla_bridge.py
index 235f53925..631ebb776 100644
--- a/jax/_src/xla_bridge.py
+++ b/jax/_src/xla_bridge.py
@@ -102,15 +102,11 @@ _CPU_ENABLE_GLOO_COLLECTIVES = config.bool_flag(
     help="Deprecated, please use jax_cpu_collectives_implementation instead.",
 )
 
-CPU_COLLECTIVES_IMPLEMENTATIONS = ["none", "gloo", "mpi"]
-CPU_COLLECTIVES_IMPLEMENTATION = config.enum_flag(
-    name="jax_cpu_collectives_implementation",
-    default="none",
-    enum_values=CPU_COLLECTIVES_IMPLEMENTATIONS,
-    help=(
-        "Cross-process collective implementation used on CPU. Must be one of"
-        f" {CPU_COLLECTIVES_IMPLEMENTATIONS}"
-    ),
+_CPU_COLLECTIVES_IMPLEMENTATION = config.string_flag(
+    name='jax_cpu_collectives_implementation',
+    default='none',
+    help='Cross-process collective implementation used on CPU. Either "none", '
+         '"gloo" or "mpi"'
 )
 
 # TODO(yueshengys): turn default back to True after resolving memory increase
@@ -135,6 +131,7 @@ _at_fork_handler_installed = False
 # Backends
 
 
+
 def tpu_client_timer_callback(timer_secs: float) -> xla_client.Client | None:
   def _log_warning():
     warnings.warn(
@@ -224,47 +221,30 @@ def register_backend_factory(name: str, factory: BackendFactory, *,
     _topology_factories[name] = make_topology
 
 
-def make_cpu_client(
-    collectives: xla_client._xla.CpuCollectives | None = None,
-) -> xla_client.Client:
-  """Creates a CPU client with the requested collectives implementation.
-
-  The implementation of CPU collectives used by the client is determined by the
-  flag `--jax_cpu_collectives_implementation` - unless `collectives` is
-  provided, in which case the flag is overridden and `collectives` is used.
-
-  Args:
-    collectives: An optional CPU collectives implementation, used by the client
-      if provided.
-
-  Raises:
-    RuntimeError: If `--jax_cpu_collectives_implementation` is unknown.
-
-  Returns:
-    The created CPU client.
-  """
-  if collectives is None:
-    collectives_impl = CPU_COLLECTIVES_IMPLEMENTATION.value
-    if _CPU_ENABLE_GLOO_COLLECTIVES.value:
-        collectives_impl = 'gloo'
-        warnings.warn('Setting `jax_cpu_enable_gloo_collectives` is '
-                      'deprecated. Please use `jax.config.update('
-                      '"jax_cpu_collectives_implementation", "gloo")` instead.',
-                      DeprecationWarning,
-                      )
-    if collectives_impl == 'gloo':
-      collectives = xla_client._xla.make_gloo_tcp_collectives(
-        distributed_client=distributed.global_state.client,
-      )
-    elif collectives_impl == 'mpi':
-      collectives = xla_client._xla.make_mpi_collectives()
-      collectives.Init()
-      atexit.register(collectives.Finalize)
-    elif collectives_impl != 'none':
-      raise RuntimeError(f"Unknown collectives implementation "
-                        f"{collectives_impl}. Available implementations are "
-                        f"{CPU_COLLECTIVES_IMPLEMENTATIONS}.")
+def make_cpu_client() -> xla_client.Client:
+  collectives: xla_client._xla.CpuCollectives | None = None
 
+  collectives_impl = _CPU_COLLECTIVES_IMPLEMENTATION.value
+  if _CPU_ENABLE_GLOO_COLLECTIVES.value:
+      collectives_impl = 'gloo'
+      warnings.warn('Setting `jax_cpu_enable_gloo_collectives` is deprecated. '
+                    'Please use `jax.config.update('
+                    '"jax_cpu_collectives_implementation", "gloo")` instead.',
+                    DeprecationWarning,
+                    )
+  if collectives_impl == 'gloo':
+    collectives = xla_client._xla.make_gloo_tcp_collectives(
+      distributed_client=distributed.global_state.client,
+    )
+  elif collectives_impl == 'mpi':
+    collectives = xla_client._xla.make_mpi_collectives()
+    collectives.Init()
+    atexit.register(collectives.Finalize)
+  elif collectives_impl != 'none':
+    collectives_impls = ['none', 'gloo', 'mpi']
+    raise RuntimeError(f"Unknown collectives implementation "
+                       f"{collectives_impl}. Available implementations are "
+                       f"{collectives_impls}.")
   return xla_client.make_cpu_client(
     asynchronous=_CPU_ENABLE_ASYNC_DISPATCH.value,
     distributed_client=distributed.global_state.client,
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
index 11ae4d268..ff7d1e059 100644
--- a/jax/experimental/mosaic/gpu/fragmented_array.py
+++ b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -268,9 +268,6 @@ class FragmentedArray:
     else:
       raise NotImplementedError(self.mlir_dtype)
 
-  def __radd__(self, other):
-    return self + other
-
   def __mul__(self, other):
     if ir.FloatType.isinstance(self.mlir_dtype):
       return self._pointwise(arith.mulf, other)
@@ -279,29 +276,16 @@ class FragmentedArray:
     else:
       raise NotImplementedError(self.mlir_dtype)
 
-  def __rmul__(self, other):
-    return self * other
-
   def __sub__(self, other):
     if not ir.FloatType.isinstance(self.mlir_dtype):
       raise NotImplementedError
     return self._pointwise(arith.subf, other)
 
-  def __rsub__(self, other):
-    if not ir.FloatType.isinstance(self.mlir_dtype):
-      raise NotImplementedError
-    return self._pointwise(lambda s, o: arith.subf(o, s), other)
-
   def __truediv__(self, other):
     if not ir.FloatType.isinstance(self.mlir_dtype):
       raise NotImplementedError
     return self._pointwise(arith.divf, other)
 
-  def __rtruediv__(self, other):
-    if not ir.FloatType.isinstance(self.mlir_dtype):
-      raise NotImplementedError
-    return self._pointwise(lambda s, o: arith.divf(o, s), other)
-
   def max(self, other):
     if not ir.FloatType.isinstance(self.mlir_dtype):
       raise NotImplementedError
@@ -583,7 +567,7 @@ class FragmentedArray:
     bw = mgpu.bytewidth(dtype)
     m, n = self.shape
     assert m % 64 == 0  # This is implied by the layout.
-    cols_per_tile = swizzle // bw
+    cols_per_tile = 128 // bw
     expected_shape = [m // 64, n // cols_per_tile, 64, cols_per_tile]
     if ir.MemRefType(ref.type).shape != expected_shape:
       raise ValueError(ref.type, (m, n))
@@ -596,7 +580,7 @@ class FragmentedArray:
     dtype = ref_ty.element_type
     bw = mgpu.bytewidth(dtype)
     m_tiles, n_tiles, m_tile_size, n_tile_size = ref_ty.shape
-    if m_tile_size != 64 or n_tile_size != (swizzle // bw):
+    if m_tile_size != 64 or n_tile_size != (128 // bw):
       raise ValueError
     m, n = m_tiles * m_tile_size, n_tiles * n_tile_size
     assert m % 64 == 0  # This is implied by the layout.
@@ -611,14 +595,13 @@ class FragmentedArray:
 
   @staticmethod
   def transfer_tiled(shape, dtype, swizzle: int | None):
-    # TODO(apaszke): We could use ldmatrix/stmatrix for 16-bit types.
     bw = mgpu.bytewidth(dtype)
     m, n = shape
-    cols_per_tile = swizzle // bw
-    if n % cols_per_tile != 0:
+    if n % 32 != 0:
       raise NotImplementedError
-    if swizzle not in {32, 64, 128}:
-      raise NotImplementedError("Only swizzled stores supported")
+    cols_per_tile = 128 // bw
+    if swizzle != 128:
+      raise NotImplementedError("Only 128B swizzle supported")
 
     c = arith.ConstantOp.create_index
     tidx = arith.remui(gpu.thread_id(gpu.Dimension.x), c(WARPGROUP_SIZE))
@@ -626,66 +609,46 @@ class FragmentedArray:
     warp_id = arith.divui(tidx, c(32))  # {0, 1, 2, 3}
     sub_row_base = arith.divui(lane_id, c(4))  # {0, 1, ..., 7}
     if bw > 2:  # Stagger is only necessary for values larger than 16bit.
-      # We split the rows into two groups (left/right) and change the order in
-      # which they perform accesses to avoid bank conflicts.
-      # It seems that the STS.64 is 2x faster (and the hardware reports no
-      # conflicts) when the conflicts are split between half-warps, as
-      # opposed to having them within the half-warp. This requires a
-      # little more work for the selects, but is ultimately worth it.
-      match swizzle:
-        case 128:
-          is_stagger_left = arith.cmpi(
-              arith.CmpIPredicate.eq, arith.remui(sub_row_base, c(2)), c(0)
-          )
-        case 64:
-          is_stagger_left = arith.cmpi(
-              arith.CmpIPredicate.eq,
-              arith.remui(arith.divui(sub_row_base, c(2)), c(2)),
-              c(0),
-          )
-        case 32:
-          # 32-byte tiles of 4-byte types have only 8 columns so there is no way
-          # to stagger the memory accesses within a single tile. We could do it
-          # across tiles, but that would be a completely different scheme.
-          raise NotImplementedError
-        case _:
-          raise AssertionError(swizzle)
-      stagger_amount = swizzle // 64
+      is_even_row = arith.cmpi(
+          arith.CmpIPredicate.eq, arith.remui(sub_row_base, c(2)), c(0)
+      )
     else:
       # We rely on canonicalization to clean up the selects.
       i1 = ir.IntegerType.get_signless(1)
-      is_stagger_left = arith.constant(i1, ir.BoolAttr.get(True))
-      stagger_amount = 0
+      is_even_row = arith.constant(i1, ir.BoolAttr.get(True))
     row_base = arith.addi(sub_row_base, arith.muli(warp_id, c(16)))
     col_base = arith.muli(arith.remui(lane_id, c(4)), c(2))  # {0, 2, 4, 6}
     # The swizzle pattern is constant for a given thread.
-    col_swizzle_bits = arith.muli(
-        arith.divui(sub_row_base, c(128 // swizzle)), c(16 // bw),
-    )
+    col_swizzle_bits = arith.muli(sub_row_base, c(16 // bw))
     for row_group in range(m // 64):
       for col_group in range(n // cols_per_tile):
         for row_subidx in range(2):
           row = arith.addi(row_base, c(row_subidx * 8))
           for col_subidx in range(cols_per_tile // 8):
-            col_subidx_left = col_subidx
-            col_subidx_right = col_subidx ^ stagger_amount
+            # We stagger the even and odd rows a little to avoid bank conflicts.
+            # It seems that the STS.64 is 2x faster (and the hardware reports no
+            # conflicts) when the conflicts are split between half-warps, as
+            # opposed to having them within the half-warp. This requires a
+            # little more work for the selects, but is ultimately worth it.
+            col_subidx_even = col_subidx
+            col_subidx_odd = col_subidx ^ 2
             col_off = arith.select(
-                is_stagger_left, c(col_subidx_left * 8), c(col_subidx_right * 8)
+                is_even_row, c(col_subidx_even * 8), c(col_subidx_odd * 8)
             )
             col = arith.addi(col_base, col_off)
             col = arith.xori(col, col_swizzle_bits)
-            reg_idx_left = col_subidx_left + col_group * (cols_per_tile // 8)
-            reg_idx_right = col_subidx_right + col_group * (cols_per_tile // 8)
-            left_idx = row_group, reg_idx_left, row_subidx, 0
-            right_idx = row_group, reg_idx_right, row_subidx, 0
+            reg_idx_even = col_subidx_even + col_group * (cols_per_tile // 8)
+            reg_idx_odd = col_subidx_odd + col_group * (cols_per_tile // 8)
+            even_idx = row_group, reg_idx_even, row_subidx, 0
+            odd_idx = row_group, reg_idx_odd, row_subidx, 0
             idx = c(row_group), c(col_group), row, col
-            def get_register(regs, left_idx=left_idx, right_idx=right_idx):
-              value_left = regs[left_idx]
-              value_right = regs[right_idx]
-              return arith.select(is_stagger_left, value_left, value_right)
-            def update_registers(regs, new, left_idx=left_idx, right_idx=right_idx):
-              regs[left_idx] = arith.select(is_stagger_left, new, regs[left_idx])
-              regs[right_idx] = arith.select(is_stagger_left, regs[right_idx], new)
+            def get_register(regs, even_idx=even_idx, odd_idx=odd_idx):
+              value_even = regs[even_idx]
+              value_odd = regs[odd_idx]
+              return arith.select(is_even_row, value_even, value_odd)
+            def update_registers(regs, new, even_idx=even_idx, odd_idx=odd_idx):
+              regs[even_idx] = arith.select(is_even_row, new, regs[even_idx])
+              regs[odd_idx] = arith.select(is_even_row, regs[odd_idx], new)
             yield get_register, update_registers, idx
 
   def tree_flatten(self):
diff --git a/jax/experimental/mosaic/gpu/profiler.py b/jax/experimental/mosaic/gpu/profiler.py
index 911783b55..4ad0e6117 100644
--- a/jax/experimental/mosaic/gpu/profiler.py
+++ b/jax/experimental/mosaic/gpu/profiler.py
@@ -140,7 +140,7 @@ class ProfilerSpec:
     return self.smem_i32_elements(block) * bytes_per_entry
 
   def intern_name(self, name: str) -> int:
-    if (name_id := self.interned_names.get(name, None)) is not None:
+    if name_id := self.interned_names.get(name, None):
       return name_id
     name_id = self.interned_names[name] = len(self.interned_names)
     if name_id & self.EXIT:
diff --git a/jax/extend/ffi.py b/jax/extend/ffi.py
index 3a26030c1..5862d0ece 100644
--- a/jax/extend/ffi.py
+++ b/jax/extend/ffi.py
@@ -20,5 +20,4 @@ from jax._src.extend.ffi import (
     ffi_lowering as ffi_lowering,
     include_dir as include_dir,
     pycapsule as pycapsule,
-    register_ffi_target as register_ffi_target,
 )
diff --git a/jaxlib/mosaic/dialect/tpu/layout.cc b/jaxlib/mosaic/dialect/tpu/layout.cc
index 1a09abdaf..7be9c85b2 100644
--- a/jaxlib/mosaic/dialect/tpu/layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/layout.cc
@@ -25,7 +25,6 @@ limitations under the License.
 #include <ostream>
 #include <string>
 #include <tuple>
-#include <utility>
 
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
@@ -446,25 +445,28 @@ SmallVector<int64_t> VectorLayout::implicitShape(
   return implicit_shape;
 }
 
-SmallVector<int64_t> VectorLayout::tileArrayShape(
-    const bool src_is_implicit, const bool res_is_implicit,
-    SmallVector<int64_t>&& src_shape,
+SmallVector<int64_t> VectorLayout::tileArrayImplicitShape(
+    const ArrayRef<int64_t> shape,
     const std::array<int64_t, 2> target_shape) const {
   const std::array<int64_t, 2> vreg_slice = vregSlice(target_shape);
-  if (!src_is_implicit) {
-    CHECK_GE(src_shape.size(), layout_rank());
-    insertImplicit<int64_t>(src_shape, 1);
-  }
-  int64_t& second_minor = *(src_shape.end() - 2);
-  int64_t& minor = *(src_shape.end() - 1);
-  second_minor =
-      llvm::divideCeil(offsets_[0].value_or(0) + second_minor, vreg_slice[0]);
-  minor = llvm::divideCeil(offsets_[1].value_or(0) + minor, vreg_slice[1]);
-  if (!res_is_implicit) {
-    CHECK_GE(src_shape.size(), 2);
-    eraseImplicit(src_shape);
-  }
-  return std::move(src_shape);
+  SmallVector<int64_t> tiles_shape = implicitShape(shape);
+  tiles_shape[tiles_shape.size() - 2] = llvm::divideCeil(
+      offsets_[0].value_or(0) + tiles_shape[tiles_shape.size() - 2],
+      vreg_slice[0]);
+  tiles_shape[tiles_shape.size() - 1] = llvm::divideCeil(
+      offsets_[1].value_or(0) + tiles_shape[tiles_shape.size() - 1],
+      vreg_slice[1]);
+  return tiles_shape;
+}
+
+SmallVector<int64_t> VectorLayout::tileArrayShape(
+    const ArrayRef<int64_t> shape,
+    const std::array<int64_t, 2> target_shape) const {
+  SmallVector<int64_t> tiles_shape =
+      tileArrayImplicitShape(shape, target_shape);
+  // Remove the implicit dimension --- it's always of size 1.
+  eraseImplicit(tiles_shape);
+  return tiles_shape;
 }
 
 std::unique_ptr<VRegDataBounds> VectorLayout::tileDataBounds(
diff --git a/jaxlib/mosaic/dialect/tpu/layout.h b/jaxlib/mosaic/dialect/tpu/layout.h
index 68dc01d84..8e7c46d8e 100644
--- a/jaxlib/mosaic/dialect/tpu/layout.h
+++ b/jaxlib/mosaic/dialect/tpu/layout.h
@@ -235,8 +235,13 @@ class VectorLayout {
     CHECK(llvm::has_single_bit<unsigned>(bitwidth_) && bitwidth_ <= 32);
   }
 
-  static int num_implicit_dims(const ImplicitDim implicit_dim) {
-    switch (implicit_dim) {
+  int8_t bitwidth() const { return bitwidth_; }
+  const LayoutOffsets &offsets() const { return offsets_; }
+  const std::array<int64_t, 2> &tiling() const { return tiling_; }
+  ImplicitDim implicit_dim() const { return implicit_dim_; }
+  int packing() const { return 32 / bitwidth_; }
+  int num_implicit_dims() const {
+    switch (implicit_dim_) {
       case ImplicitDim::kNone:
         return 0;
       case ImplicitDim::kMinor:
@@ -244,19 +249,8 @@ class VectorLayout {
         return 1;
     }
   }
-
-  // The number of non-implicit dimensions that are tiled.
-  static int layout_rank(const ImplicitDim implicit_dim) {
-    return 2 - num_implicit_dims(implicit_dim);
-  }
-
-  int8_t bitwidth() const { return bitwidth_; }
-  const LayoutOffsets &offsets() const { return offsets_; }
-  const std::array<int64_t, 2> &tiling() const { return tiling_; }
-  ImplicitDim implicit_dim() const { return implicit_dim_; }
-  int packing() const { return 32 / bitwidth_; }
-  int num_implicit_dims() const { return num_implicit_dims(implicit_dim_); }
-  int layout_rank() const { return layout_rank(implicit_dim_); }
+  // The number of minormost dimensions tiled by this layout.
+  int layout_rank() const { return 2 - num_implicit_dims(); }
 
   bool operator==(const VectorLayout &other) const;
   bool operator!=(const VectorLayout &other) const {
@@ -314,20 +308,6 @@ class VectorLayout {
     }
   }
 
-  static std::array<int64_t, 2> getImplicitTiledDims(
-      const ImplicitDim implicit_dim, const ArrayRef<int64_t> arr,
-      const int64_t implicit_value) {
-    CHECK_GE(arr.size(), layout_rank(implicit_dim));
-    switch (implicit_dim) {
-      case ImplicitDim::kNone:
-        return {*(arr.end() - 2), *(arr.end() - 1)};
-      case ImplicitDim::kMinor:
-        return {*(arr.end() - 1), implicit_value};
-      case ImplicitDim::kSecondMinor:
-        return {implicit_value, *(arr.end() - 1)};
-    }
-  }
-
   // Returns the value of the tiled (2 minormost) dimensions of the given array
   // with implicit dims inserted.
   //
@@ -338,11 +318,23 @@ class VectorLayout {
   //   return {*(vec.end() - 2), *(vec.end() - 1)};
   std::array<int64_t, 2> getImplicitTiledDims(
       const ArrayRef<int64_t> arr, const int64_t implicit_value) const {
-    return getImplicitTiledDims(implicit_dim_, arr, implicit_value);
+    CHECK_GE(arr.size(), layout_rank());
+    switch (implicit_dim_) {
+      case ImplicitDim::kNone:
+        return {*(arr.end() - 2), *(arr.end() - 1)};
+      case ImplicitDim::kMinor:
+        return {*(arr.end() - 1), implicit_value};
+      case ImplicitDim::kSecondMinor:
+        return {implicit_value, *(arr.end() - 1)};
+    }
   }
 
   SmallVector<int64_t> implicitShape(ArrayRef<int64_t> shape) const;
 
+  SmallVector<int64_t> tileArrayImplicitShape(
+      ArrayRef<int64_t> shape, std::array<int64_t, 2> target_shape) const;
+
+ public:
   // Returns the shape of ndarray of vregs needed to represent a value.
   //
   // All but the last two dimensions are unrolled over vregs. In the last two
@@ -353,31 +345,9 @@ class VectorLayout {
   // minimize the number of vregs.
   //
   // Args:
-  //   src_is_implicit: If true, the input shape already has implicit dimensions
-  //     inserted.
-  //   res_is_implicit: If true, the output shape will have implicit dimensions
-  //     inserted.
-  //   shape: The shape of the full vector this layout applies to, possibly
-  //     with implicit dimensions inserted.
-  SmallVector<int64_t> tileArrayShape(
-      bool src_is_implicit, bool res_is_implicit, ArrayRef<int64_t> shape,
-      std::array<int64_t, 2> target_shape) const {
-    return tileArrayShape(src_is_implicit, res_is_implicit,
-                          SmallVector<int64_t>(shape), target_shape);
-  }
-  SmallVector<int64_t> tileArrayShape(
-      bool src_is_implicit, bool res_is_implicit, SmallVector<int64_t> &&shape,
-      std::array<int64_t, 2> target_shape) const;
-
-  SmallVector<int64_t> tileArrayImplicitShape(
-      ArrayRef<int64_t> shape, std::array<int64_t, 2> target_shape) const {
-    return tileArrayShape(false, true, shape, target_shape);
-  }
-
+  //   shape: The shape of the full vector this layout applies to.
   SmallVector<int64_t> tileArrayShape(
-      ArrayRef<int64_t> shape, std::array<int64_t, 2> target_shape) const {
-    return tileArrayShape(false, false, shape, target_shape);
-  }
+      ArrayRef<int64_t> shape, std::array<int64_t, 2> target_shape) const;
 
   // Returns the bounds of the given tile that hold useful data.
   //
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
index 2580f1324..4f11e7aef 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -2300,7 +2300,8 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
     }
     // Convert dynamic shift to log(bound) static ops.
     int roll_by = 1;
-    while (roll_by < bound) {
+    Value cst_1 = mlirI32Const(1);
+    while (bound > 0) {
       auto new_result = rotate(
           result,
           mlirI32Const(tiling_dim >= 0 ? roll_by * ctx.target_shape[tiling_dim]
@@ -2309,14 +2310,15 @@ LogicalResult rotate_rule_impl(RewriteContext &ctx, OpTy op, Value amount,
       auto mask = builder.create<arith::CmpIOp>(
           arith::CmpIPredicate::ne,
           builder.create<vector::BroadcastOp>(
-              i32_vreg,
-              builder.create<arith::AndIOp>(vreg_shift, mlirI32Const(roll_by))),
+              i32_vreg, builder.create<arith::AndIOp>(vreg_shift, cst_1)),
           builder.create<arith::ConstantOp>(
               DenseElementsAttr::get(i32_vreg, builder.getI32IntegerAttr(0))));
       result.Each([&](absl::Span<const int64_t> idxs, Value *v) {
         *v = builder.create<arith::SelectOp>(mask, new_result(idxs), *v);
       });
       roll_by *= 2;
+      bound /= 2;
+      vreg_shift = divI(vreg_shift, 2);
     }
     return result;
   };
@@ -3829,44 +3831,55 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
     no_op = true;
   } else if (  // Fold or unfold sublane dim, but keeping a whole number of
                // vregs.
+      layout_in.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+      layout_out.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
       layout_in.offsets()[0] == 0 &&
       layout_in.offsets() == layout_out.offsets() &&
       layout_in.tiling() == layout_out.tiling() &&
-      dst_tiled_dims[1] == src_tiled_dims[1] &&
-      dst_tiled_dims[0] % dst_vreg_slice[0] == 0 &&
-      src_tiled_dims[0] % src_vreg_slice[0] == 0) {
+      *(dst_shape.end() - 1) == *(src_shape.end() - 1) &&
+      *(dst_shape.end() - 2) % dst_vreg_slice[0] == 0 &&
+      *(src_shape.end() - 2) % src_vreg_slice[0] == 0) {
     no_op = true;
-  } else if (layout_in.offsets() == layout_out.offsets() &&
+  } else if (layout_in.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+             layout_out.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+             layout_in.offsets() == layout_out.offsets() &&
              layout_in.offsets() == LayoutOffsets{0, 0} &&
              layout_in.tiling()[0] == 1 &&
              layout_out.hasNativeTiling(ctx.target_shape) &&
-             dst_tiled_dims[1] == dst_vreg_slice[1] &&
-             dst_tiled_dims[0] % dst_vreg_slice[0] == 0 &&
-             src_tiled_dims[1] % src_vreg_slice[1] == 0) {
+             *(dst_shape.end() - 1) == dst_vreg_slice[1] &&
+             *(dst_shape.end() - 2) % dst_vreg_slice[0] == 0 &&
+             *(src_shape.end() - 1) % src_vreg_slice[1] == 0) {
     // Shapecast (..., m * 128 * packing) -> (..., 128).
     no_op = true;
-  } else if (layout_in.offsets() == LayoutOffsets{0, 0} &&
+  } else if (layout_in.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+             layout_out.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
+             layout_in.offsets() == LayoutOffsets{0, 0} &&
              layout_out.offsets() == LayoutOffsets{0, 0} &&
              layout_in.hasNativeTiling(ctx.target_shape) &&
              layout_out.tiling()[0] == 1 &&
-             src_tiled_dims[1] == src_vreg_slice[1] &&
-             src_tiled_dims[0] % src_vreg_slice[0] == 0 &&
-             dst_tiled_dims[1] % dst_vreg_slice[1] == 0) {
+             *(src_shape.end() - 1) == src_vreg_slice[1] &&
+             *(src_shape.end() - 2) % src_vreg_slice[0] == 0 &&
+             *(dst_shape.end() - 1) % dst_vreg_slice[1] == 0) {
     // Shapecast (..., 128) -> (..., m * 128 * packing).
     no_op = true;
   }
   FAILUREOR_ASSIGN_OR_RETURN(
       xla::Array<Value> src_vregs,
       disassemble(builder, layout_in, shape_cast_op.getSource(),
-                  ctx.target_shape, /*use_implicit_shape=*/true));
+                  ctx.target_shape));
   auto getDstVregs = [&]() -> FailureOr<xla::Array<Value>> {
     if (no_op) {
       xla::Array<Value> dst_vregs_local = src_vregs;
       dst_vregs_local.Reshape(
-          layout_out.tileArrayImplicitShape(dst_shape, ctx.target_shape));
+          layout_out.tileArrayShape(dst_shape, ctx.target_shape));
       return dst_vregs_local;
-    } else if (dst_tiled_dims == std::array<int64_t, 2>{src_tiled_dims[1], 1} &&
+    } else if (dst_shape.take_back(2) ==
+                   ArrayRef<int64_t>{src_shape.back(), 1} &&
                layout_in.bitwidth() == 32 &&
+               (layout_in.implicit_dim() == VectorLayout::ImplicitDim::kNone ||
+                layout_in.implicit_dim() ==
+                    VectorLayout::ImplicitDim::kSecondMinor) &&
+               layout_out.implicit_dim() == VectorLayout::ImplicitDim::kNone &&
                layout_in.hasNativeTiling(ctx.target_shape) &&
                layout_in.tiling() == layout_out.tiling() &&
                layout_in.offsets()[0].value_or(0) == 0 &&
@@ -3875,11 +3888,21 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
                // replicated result
     ) {
       // First, insert the new singleton lane dimension.
-      SmallVector<int64_t> s = layout_in.implicitShape(src_shape);
+      SmallVector<int64_t> s(src_shape);
       s.push_back(1);
-      xla::Array<Value> dst_vregs_local(layout_out.tileArrayShape(
-          /*src_is_implicit=*/true, /*res_is_implicit=*/true, std::move(s),
-          ctx.target_shape));
+      xla::Array<Value> dst_vregs_local(
+          layout_out.tileArrayShape(s, ctx.target_shape));
+      if (layout_in.implicit_dim() == VectorLayout::ImplicitDim::kSecondMinor) {
+        // Make the sublane dimension explicit.
+        SmallVector<int64_t> new_src_vregs_shape(
+            toArrayRef(src_vregs.dimensions()));
+        new_src_vregs_shape.insert(new_src_vregs_shape.end() - 1, 1);
+        src_vregs.Reshape(new_src_vregs_shape);
+        SmallVector<int64_t> new_dst_vregs_shape(
+            toArrayRef(dst_vregs_local.dimensions()));
+        new_dst_vregs_shape.insert(new_dst_vregs_shape.end() - 2, 1);
+        dst_vregs_local.Reshape(new_dst_vregs_shape);
+      }
       TPU_ASSERT_EQ_OP(dst_vregs_local.dimensions().back(),
                        1);  // We're inserting a singleton dimension
       dst_vregs_local.Each(
@@ -3920,9 +3943,9 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
           *dst_vreg = dst_vregs_local(first_row_idx);
         });
       }
-      // Now, reshape the major axes of the vreg array.
+      // Now, permute the major axes of the vreg array.
       dst_vregs_local.Reshape(
-          layout_out.tileArrayImplicitShape(dst_shape, ctx.target_shape));
+          layout_out.tileArrayShape(dst_shape, ctx.target_shape));
       return dst_vregs_local;
     } else {
       return shape_cast_op.emitOpError(
@@ -3931,9 +3954,8 @@ LogicalResult vector_shape_cast_rule(RewriteContext &ctx, Operation &op,
     }
   };
   FAILUREOR_ASSIGN_OR_RETURN(const xla::Array<Value> dst_vregs, getDstVregs());
-  shape_cast_op->replaceAllUsesWith(assemble(builder, dst_ty, layout_out,
-                                             dst_vregs, ctx.target_shape,
-                                             /*use_implicit_shape=*/true));
+  shape_cast_op->replaceAllUsesWith(
+      assemble(builder, dst_ty, layout_out, dst_vregs, ctx.target_shape));
   shape_cast_op->erase();
   return success();
 }
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
index 3fd3a0bea..1a5639990 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -1341,66 +1341,39 @@ class VectorLayoutInferer {
   LogicalResult infer(vector::ShapeCastOp op) {
     auto src_ty = op.getSourceVectorType();
     auto src_shape = src_ty.getShape();
+    int64_t src_rank = src_ty.getRank();
     auto res_ty = op.getResultVectorType();
     auto res_shape = res_ty.getShape();
+    int64_t res_rank = res_ty.getRank();
     auto some_src_layout = getLayout(op.getSource());
     TPU_CHECK_OP(some_src_layout, "missing vector layout");
     auto layout = *some_src_layout;
     const unsigned bitwidth = src_ty.getElementTypeBitWidth();
-    const std::array<int64_t, 2> native_tiling = nativeTiling(bitwidth);
-    const std::array<int64_t, 2> src_tiled_ishape =
-        layout.getImplicitTiledDims(src_shape, 1);
     const std::array<int64_t, 2> vreg_slice = layout.vregSlice(target_shape_);
-
-    // TODO(tlongeri): Be smarter about trying implicit dims. We should probably
-    // only add them when folding dimensions, and remove them when unfolding.
-    SmallVector<ImplicitDim, 3> candidate_implicit_dims;
-    if (res_shape.size() >= 2) {
-      candidate_implicit_dims.push_back(ImplicitDim::kNone);
-    }
-    if (!res_shape.empty()) {
-      candidate_implicit_dims.push_back(ImplicitDim::kSecondMinor);
-      candidate_implicit_dims.push_back(ImplicitDim::kMinor);
-    }
-    // TODO(b/340625465): Add case with both implicit dims once we support it.
-
-    // See if we can get implicit tiled dimensions to match. This is always a
-    // no-op.
-    for (const ImplicitDim implicit_dim : candidate_implicit_dims) {
-      const std::array<int64_t, 2> res_tiled_ishape =
-          VectorLayout::getImplicitTiledDims(implicit_dim, res_shape, 1);
-      if (src_tiled_ishape == res_tiled_ishape) {
-        // Nothing changes in the tiled dimensions
-        setLayout(op, layout,
-                  VectorLayout(layout.bitwidth(), layout.offsets(),
-                               layout.tiling(), implicit_dim));
+    if (layout.implicit_dim() == ImplicitDim::kNone) {
+      // Nothing changes in the last two dims.
+      if (res_rank >= 2 && src_shape.take_back(2) == res_shape.take_back(2)) {
+        setLayout(op, layout, layout);
         return success();
       }
-    }
-
-    // See if we can do sublane or lane (un)folding.
-    for (const ImplicitDim implicit_dim : candidate_implicit_dims) {
-      const std::array<int64_t, 2> res_tiled_ishape =
-          VectorLayout::getImplicitTiledDims(implicit_dim, res_shape, 1);
-      // Sublane (un)folding.
-      if (src_tiled_ishape[1] == res_tiled_ishape[1] &&
-          src_tiled_ishape[0] % vreg_slice[0] == 0 &&
-          res_tiled_ishape[0] % vreg_slice[0] == 0) {
+      // Sublane (un)tiling.
+      if (res_rank >= 2 && *(src_shape.end() - 1) == *(res_shape.end() - 1) &&
+          *(src_shape.end() - 2) % vreg_slice[0] == 0 &&
+          *(res_shape.end() - 2) % vreg_slice[0] == 0) {
         // TODO(b/343808585): We shouldn't force second minor offset to 0 when
         //                    unfolding, it's still a no-op, but we need to add
         //                    support in apply-vector-layout.
-        const LayoutOffsets offsets = {0, layout.offsets()[1]};
-        setLayout(op,
-                  VectorLayout(layout.bitwidth(), offsets, layout.tiling(),
-                               layout.implicit_dim()),
-                  VectorLayout(layout.bitwidth(), offsets, layout.tiling(),
-                               implicit_dim));
+        layout = VectorLayout(layout.bitwidth(), {0, layout.offsets()[1]},
+                              layout.tiling(), layout.implicit_dim());
+        setLayout(op, layout, layout);
         return success();
       }
-      // Lane (un)folding.
-      if (src_tiled_ishape[1] != res_tiled_ishape[1] &&
-          src_tiled_ishape[1] % layout.tiling()[1] == 0 &&
-          res_tiled_ishape[1] % layout.tiling()[1] == 0) {
+      const auto native_tiling = nativeTiling(bitwidth);
+      // Lane (un)tiling.
+      if (src_ty.getDimSize(src_ty.getRank() - 1) !=
+              res_shape[res_shape.size() - 1] &&
+          src_ty.getDimSize(src_ty.getRank() - 1) % layout.tiling()[1] == 0 &&
+          res_shape[res_shape.size() - 1] % layout.tiling()[1] == 0) {
         const int packing = kNativeBitwidth / bitwidth;
         const auto elements_per_vreg = native_tiling[0] * native_tiling[1];
         // When we shapecast from input shape
@@ -1408,17 +1381,17 @@ class VectorLayoutInferer {
         // (..., target_shape_[1]), the reshape becomes no-op when input is
         // densely packed with tiling (1, target_shape_[1] * packing) and output
         // has the native tiling.
-        if (res_tiled_ishape[1] == target_shape_[1] &&
-            res_tiled_ishape[0] % native_tiling[0] == 0 &&
-            src_tiled_ishape[1] % elements_per_vreg == 0) {
+        if (*(res_shape.end() - 1) == target_shape_[1] &&
+            *(res_shape.end() - 2) % native_tiling[0] == 0 &&
+            *(src_shape.end() - 1) % elements_per_vreg == 0) {
           // Inferring in_layout to have tiling (1, 128 * packing) triggers any
           // necessary relayout before shapecast.
-          setLayout(op,
-                    VectorLayout(layout.bitwidth(), {0, 0},
-                                 {1, target_shape_[1] * packing},
-                                 layout.implicit_dim()),
-                    VectorLayout(layout.bitwidth(), {0, 0}, native_tiling,
-                                 implicit_dim));
+          setLayout(
+              op,
+              VectorLayout(layout.bitwidth(), {0, 0},
+                           {1, target_shape_[1] * packing}, ImplicitDim::kNone),
+              VectorLayout(layout.bitwidth(), {0, 0}, native_tiling,
+                           ImplicitDim::kNone));
           return success();
         }
 
@@ -1426,32 +1399,98 @@ class VectorLayoutInferer {
         // shape (..., m * target_shape_[1] * packing), the reshape becomes
         // no-op when input has the native tiling and output is densely packed
         // with tiling (1, target_shape_[1] * packing).
-        if (src_tiled_ishape[1] == target_shape_[1] &&
-            src_tiled_ishape[0] % native_tiling[0] == 0 &&
-            res_tiled_ishape[1] % elements_per_vreg == 0) {
-          setLayout(
-              op,
-              VectorLayout(layout.bitwidth(), {0, 0}, native_tiling,
-                           layout.implicit_dim()),
-              VectorLayout(layout.bitwidth(), {0, 0},
-                           {1, target_shape_[1] * packing}, implicit_dim));
+        if (*(src_shape.end() - 1) == target_shape_[1] &&
+            *(src_shape.end() - 2) % native_tiling[0] == 0 &&
+            *(res_shape.end() - 1) % elements_per_vreg == 0) {
+          setLayout(op,
+                    VectorLayout(layout.bitwidth(), {0, 0}, native_tiling,
+                                 ImplicitDim::kNone),
+                    VectorLayout(layout.bitwidth(), {0, 0},
+                                 {1, target_shape_[1] * packing},
+                                 ImplicitDim::kNone));
           return success();
         }
+
+        // TODO(b/299253805): support shapecast along lane for other cases.
+        op.emitOpError("unsupported shape cast");
+        return failure();
       }
-    }
+      if (layout.tiling() != native_tiling) {
+        layout = VectorLayout(bitwidth, layout.offsets(), native_tiling,
+                              layout.implicit_dim());
+      }
+      TPU_CHECK_OP(src_ty.getRank() >= 2,
+                   "expected 2D+ operand with 2D layout");
+      ArrayRef<int64_t> layout_shape = src_ty.getShape().take_back(2);
+      if (res_ty.getRank() >= 2) {
+        // Squeeze out the sublane dim.
+        if (layout_shape[0] == 1 &&
+            res_shape.back() == src_shape.back()) {
+          setLayout(op, layout,
+                    VectorLayout(bitwidth, layout.offsets(), layout.tiling(),
+                                 ImplicitDim::kSecondMinor));
+          return success();
+        }
+        // Insert a singleton lane dimension. The old lane dimension ends up
+        // in the sublane dimension. Other axes can be reshaped arbitrarily.
+        if (src_ty.getElementTypeBitWidth() == kNativeBitwidth &&
+            src_shape.back() == res_shape[res_shape.size() - 2] &&
+            res_shape.back() == 1) {
+          setLayout(op, layout,
+                    VectorLayout(kNativeBitwidth, {0, std::nullopt},
+                                 default_tiling_, ImplicitDim::kNone));
+          return success();
+        }
+      } else if (res_ty.getRank() == 1) {
+        // All dimensions have been folded into a single one
 
-    // Try adding a singleton innermost dim to the actual *implicit* shape.
-    if (res_shape.size() >= 2 &&
-        res_shape.take_back(2) == ArrayRef<int64_t>({src_tiled_ishape[1], 1})) {
-      TPU_CHECK_OP(bitwidth == kNativeBitwidth,
-                   "Insertion of minor dim that is not a no-op only "
-                   "supported for 32-bit types");
-      setLayout(op,
-                VectorLayout(layout.bitwidth(), layout.offsets(), native_tiling,
-                             layout.implicit_dim()),
-                VectorLayout(layout.bitwidth(), {0, std::nullopt},
-                             native_tiling, ImplicitDim::kNone));
-      return success();
+        // Squeeze all but minor dimension
+        if (res_ty.getShape().back() == layout_shape[1]) {
+          // The condition implies that everything apart from the minor
+          // dimension is 1 in the source.
+          setLayout(op, layout,
+                    VectorLayout(bitwidth, layout.offsets(), layout.tiling(),
+                                 ImplicitDim::kSecondMinor));
+          return success();
+        }
+        // Squeeze all but second minor dimension
+        if (res_ty.getShape().back() == layout_shape[0]) {
+          // The condition implies that everything apart from the second minor
+          // dimension is 1 in the source
+          setLayout(op, layout,
+                    VectorLayout(kNativeBitwidth, layout.offsets(),
+                                 layout.tiling(), ImplicitDim::kMinor));
+          return success();
+        }
+        // TODO(b/340625465): Add case where layout_shape is (1, 1) and we fold
+        //                    batch dimensions once we support 0-D layouts.
+      }
+    } else {
+      // Nothing changes in the last dim.
+      if (res_ty.getRank() >= 1 && src_shape.back() == res_shape.back()) {
+        setLayout(op, layout, layout);
+        return success();
+      }
+      // Insert a singleton innermost dim.
+      if (res_ty.getRank() == src_ty.getRank() + 1 &&
+          src_ty.getDimSize(src_rank - 1) == res_ty.getDimSize(res_rank - 2) &&
+          res_ty.getDimSize(res_rank - 1) == 1) {
+        if (layout.implicit_dim() == ImplicitDim::kMinor) {
+          setLayout(op, layout,
+                    VectorLayout(bitwidth, layout.offsets(), layout.tiling(),
+                                 ImplicitDim::kNone));
+        } else {
+          TPU_CHECK_OP(bitwidth == kNativeBitwidth,
+                       "Insertion of minor dim that is not a no-op only "
+                       "supported for 32-bit types");
+          TPU_CHECK_OP(layout.implicit_dim() == ImplicitDim::kSecondMinor,
+                       "unexpected implicit dim value");
+          setLayout(op, layout,
+                    VectorLayout(bitwidth, {0, std::nullopt}, default_tiling_,
+                                 ImplicitDim::kNone));
+        }
+        return success();
+      }
     }
     op.emitOpError("unsupported shape cast");
     return failure();
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index ce670af71..04bdc251a 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 = "e608ef43d7a4e719d40366a2f09e75ddd793862b"
-XLA_SHA256 = "af49038019d4849738f52f7f70c86d17b87d899fa673887dd5795dfd734c8597"
+XLA_COMMIT = "6695d15d8e4b58b84174317c9ebcadc4f701cc5b"
+XLA_SHA256 = "b962f43b76b5ff3283af046be44ac0965d17489d764c05abc4592ffb9aa88ad7"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/lax_metal_test.py b/tests/lax_metal_test.py
index 02fecb7b3..dab26d86c 100644
--- a/tests/lax_metal_test.py
+++ b/tests/lax_metal_test.py
@@ -578,7 +578,6 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       dtype=float_dtypes,#number_dtypes,
   )
   @jax.default_matmul_precision("float32")
-  @jax.numpy_rank_promotion('allow') # adopt PR#22316
   def testVecdot(self, lhs_batch, rhs_batch, axis_size, axis, dtype):
     # Construct vecdot-compatible shapes.
     size = min(len(lhs_batch), len(rhs_batch))
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index 6ec29a494..009560f87 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -364,28 +364,20 @@ class WGMMATest(TestCase):
     )()
     np.testing.assert_array_equal(iota, expected)
 
-  @parameterized.product(
-      dtypes=(
-          (ir.F32Type.get, jnp.float32),
-          (ir.F16Type.get, jnp.float16),
-          (partial(ir.IntegerType.get_signless, 8), jnp.int8),
-      ),
-      swizzle=(32, 64, 128),
-      num_col_tiles=(1, 2, 3),
+  @parameterized.named_parameters(
+      ("f32", ir.F32Type.get, jnp.float32),
+      ("f16", ir.F16Type.get, jnp.float16),
+      ("i8", partial(ir.IntegerType.get_signless, 8), jnp.int8),
   )
-  def test_store_tiled(self, dtypes, swizzle, num_col_tiles):
-    mlir_dtype_cls, jax_dtype = dtypes
+  def test_store_tiled(self, mlir_dtype_cls, jax_dtype):
     mlir_dtype = mlir_dtype_cls()
-    if bytewidth(mlir_dtype) > 2 and swizzle == 32:
-      self.skipTest("Not implemented")
-    col_tiling = swizzle // bytewidth(mlir_dtype)
     m = 128
-    n = col_tiling * num_col_tiles
-    tiling = (64, col_tiling)
+    n = 256
+    tiling = (64, 128 // bytewidth(mlir_dtype))
     def kernel(ctx, out, smem):
       del ctx
-      iota_tensor(m, n, mlir_dtype).store_tiled(smem, swizzle=swizzle)
-      copy(smem, out, swizzle=swizzle)
+      iota_tensor(m, n, mlir_dtype).store_tiled(smem, swizzle=128)
+      copy(smem, out, swizzle=128)
     expected = (
         np.arange(m * n, dtype=jax_dtype)
         .reshape(m // tiling[0], tiling[0], n // tiling[1], tiling[1])
diff --git a/tests/pallas/tpu_pallas_pipeline_test.py b/tests/pallas/tpu_pallas_pipeline_test.py
index 1d75949b8..43f529c6a 100644
--- a/tests/pallas/tpu_pallas_pipeline_test.py
+++ b/tests/pallas/tpu_pallas_pipeline_test.py
@@ -400,6 +400,12 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
       # Fixed Ref schedule, only really needed to prevent HBM data race in the
       # degenerate case of a trivial (single-step) inner loop.
       accum_schedule = pltpu.get_pipeline_schedule('fixed')
+      # Tweak schedule to skip copying in initial accumulator data as we zero it
+      # out anyway.
+      for k in ['prologue_copy_in', 'wait_in', 'copy_in']:
+        accum_schedule[k] = functools.partial(  # avoid cell-var-from-loop
+            lambda original_pred_fn, *a: original_pred_fn(*a) & ~is_start,
+            accum_schedule[k])
 
       # Outer loop prologue
       @pl.when(is_start)
@@ -883,8 +889,6 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
             pl.ds(pl.multiple_of(0, sharded_k), sharded_k),
         )
 
-      rhs_schedule = pltpu.get_pipeline_schedule('fixed')
-
       # Outer Loop Prologue
       @pl.when(is_start)
       @jax.named_scope('sync')
@@ -926,13 +930,10 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
               bwd_copy.wait()
 
           # deferred "prefetch"
-          # Don't prefetch when accums are being zeroed.
-          @pl.when(~is_start)
-          def _prefetch():
-            next_slot = jnp.where(phase == 0, working_slot, buffering_slot)
-            next_phase = 1 - phase
-            scheduler.prefetch(out_bref,
-                                accumulator_ref.at[next_phase, next_slot])
+          next_slot = jnp.where(phase == 0, working_slot, buffering_slot)
+          next_phase = 1 - phase
+          scheduler.prefetch(out_bref,
+                              accumulator_ref.at[next_phase, next_slot])
 
       # Prefetch next inputs on last step of our present cycle
       def prefetch(lhs_bref, rhs_bref, out_bref, scheduler):
@@ -953,11 +954,11 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
         next_phase = lax.rem(phase + 1, 2)
         scheduler.prefetch(
             lhs_bref, lhs_ref.at[get_lhs_slice(next_step, next_phase)])
-        scheduler.prefetch(rhs_bref, rhs_ref, rhs_schedule)
+        scheduler.prefetch(
+            rhs_bref, rhs_ref)
         # When the inner matmul loop consists of a single iteration, we need
         # to avoid optimistic prefetch to avoid a data race.
-        # Don't prefetch when accums are being zeroed.
-        @pl.when(~trivial_loop & ~is_start)
+        @pl.when(~trivial_loop)
         def _prefetch_accum():
           scheduler.prefetch(
               out_bref, accumulator_ref.at[next_phase, next_working_slot])
@@ -974,7 +975,6 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
           init_accumulators=outer_step == 0,
           prefetch=prefetch,
           postyeet=postyeet,
-          schedule=[None, rhs_schedule, None],
       )
 
       # Add forwards and backwards stream results together
@@ -1160,8 +1160,6 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
       def get_accum_slice(phase, slot):
         return (slot, make_ds(phase, half_m))
 
-      rhs_schedule = pltpu.get_pipeline_schedule('fixed')
-
       # Outer Loop Prologue
       @pl.when(is_start)
       @jax.named_scope('sync')
@@ -1201,15 +1199,12 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
             def _wait_prev_bwd_dma():
               bwd_copy.wait()
           # deferred "prefetch"
-          # Don't prefetch when accums are being zeroed.
-          @pl.when(~is_start)
-          def _prefetch():
-            next_working_slot = jnp.where(
-                phase == 0, working_slot, buffering_slot)
-            next_phase = 1 - phase
-            scheduler.prefetch(
-                out_bref, rs_accum_scratch_ref.at[
-                    get_accum_slice(next_phase, next_working_slot)])
+          next_working_slot = jnp.where(
+              phase == 0, working_slot, buffering_slot)
+          next_phase = 1 - phase
+          scheduler.prefetch(
+              out_bref, rs_accum_scratch_ref.at[
+                  get_accum_slice(next_phase, next_working_slot)])
 
       # Prefetch next inputs on last step of our present cycle
       def prefetch(lhs_bref, rhs_bref, out_bref, scheduler):
@@ -1230,9 +1225,8 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
         next_phase = lax.rem(phase + 1, 2)
         scheduler.prefetch(lhs_bref,
                            lhs_ref.at[get_lhs_slice(next_step, next_phase)])
-        scheduler.prefetch(rhs_bref, rhs_ref, rhs_schedule)
-        # Don't prefetch when accums are being zeroed.
-        @pl.when(~trivial_loop & ~is_start)
+        scheduler.prefetch(rhs_bref, rhs_ref)
+        @pl.when(~trivial_loop)
         def _prefetch_accumulator():
           scheduler.prefetch(
               out_bref, rs_accum_scratch_ref.at[
@@ -1250,7 +1244,6 @@ class PallasCallCollectivePipelineTest(parameterized.TestCase):
           init_accumulators=outer_step == 0,
           prefetch=prefetch,
           postyeet=postyeet,
-          schedule=[None, rhs_schedule, None],
       )
 
     kernel = pl.pallas_call(
diff --git a/tests/profiler_test.py b/tests/profiler_test.py
index b67b078ae..a261e7f12 100644
--- a/tests/profiler_test.py
+++ b/tests/profiler_test.py
@@ -22,6 +22,7 @@ import threading
 import time
 import unittest
 from absl.testing import absltest
+import pathlib
 
 import jax
 import jax.numpy as jnp
@@ -103,6 +104,26 @@ class ProfilerTest(unittest.TestCase):
         self.assertIn(b"/device:TPU", proto)
       self.assertIn(b"pxla.py", proto)
 
+  def testProgrammaticProfilingPathlib(self):
+    with tempfile.TemporaryDirectory() as tmpdir_string:
+      tmpdir = pathlib.Path(tmpdir_string)
+      try:
+        jax.profiler.start_trace(tmpdir)
+        jax.pmap(lambda x: jax.lax.psum(x + 1, 'i'), axis_name='i')(
+            jnp.ones(jax.local_device_count()))
+      finally:
+        jax.profiler.stop_trace()
+
+      proto_path = tuple(tmpdir.rglob("*.xplane.pb"))
+      self.assertEqual(len(proto_path), 1)
+      proto = proto_path[0].read_bytes()
+      # Sanity check that serialized proto contains host, device, and
+      # Python traces without deserializing.
+      self.assertIn(b"/host:CPU", proto)
+      if jtu.test_device_matches(["tpu"]):
+        self.assertIn(b"/device:TPU", proto)
+      self.assertIn(b"pxla.py", proto)
+
   def testProfilerGetFDOProfile(self):
     # Tests stop_and_get_fod_profile could run.
     try:
@@ -147,6 +168,22 @@ class ProfilerTest(unittest.TestCase):
       if jtu.test_device_matches(["tpu"]):
         self.assertIn(b"/device:TPU", proto)
 
+  def testProgrammaticProfilingContextManagerPathlib(self):
+    with tempfile.TemporaryDirectory() as tmpdir_string:
+      tmpdir = pathlib.Path(tmpdir_string)
+      with jax.profiler.trace(tmpdir):
+        jax.pmap(lambda x: jax.lax.psum(x + 1, 'i'), axis_name='i')(
+            jnp.ones(jax.local_device_count()))
+
+      proto_path = tuple(tmpdir.rglob("*.xplane.pb"))
+      self.assertEqual(len(proto_path), 1)
+      proto = proto_path[0].read_bytes()
+      # Sanity check that serialized proto contains host and device traces
+      # without deserializing.
+      self.assertIn(b"/host:CPU", proto)
+      if jtu.test_device_matches(["tpu"]):
+        self.assertIn(b"/device:TPU", proto)
+
   def testTraceAnnotation(self):
     x = 3
     with jax.profiler.TraceAnnotation("mycontext"):