Repository: google/jax
Pull request #20251 (2024-03-14 14:52:45)

Fix `axis=None, keepdims=True` for `jnp.quantile` and `jnp.median`
This PR fixes the incorrect behavior:

```bash
>>> x = np.ones([2, 3, 4])
>>> np.median(x, axis=None, keepdims=True).shape
(1, 1, 1)
>>> jnp.median(x, axis=None, keepdims=True).shape
(1,)  # mismatched!
>>> np.quantile(x, 0.5, axis=None, keepdims=True).shape
(1, 1, 1)
>>> jnp.quantile(x, 0.5, axis=None, keepdims=True).shape
(1,)  # mismatched!
```

The corresponding tests have been included

jakevdp: Good catch, thanks for the fix! One minor suggestion: rather than new specific test cases, we could add relevant cases to the old tests. For example here: https://github.com/google/jax/blob/64bd95ded50fc647d6e80a55ab1cb3baef368cf6/tests/lax_numpy_reducers_test.py#L664
and here: https://github.com/google/jax/blob/64bd95ded50fc647d6e80a55ab1cb3baef368cf6/tests/lax_numpy_reducers_test.py#L718
We could add another case that looks something like `((6, 7), None)` and that would lead to this case being covered for our runs with larger `JAX_NUM_GENERATED_CASES`. What do you think?
james77777778: > We could add another case that looks something like `((6, 7), None)` and that would lead to this case being covered for our runs with larger `JAX_NUM_GENERATED_CASES`. What do you think?

I'm not familar with the test configuration of JAX. Thanks for the suggestion.
The tests have been updated and the new tests should be included using `JAX_NUM_GENERATED_CASES=40`

Patch:
diff --git a/README.md b/README.md
index 44a6d6e26..522d20cea 100644
--- a/README.md
+++ b/README.md
@@ -17,8 +17,9 @@
 
 ## What is JAX?
 
-JAX is a Python library for accelerator-oriented array computation and program transformation,
-designed for high-performance numerical computing and large-scale machine learning.
+JAX is [Autograd](https://github.com/hips/autograd) and [XLA](https://www.tensorflow.org/xla),
+brought together for high-performance numerical computing, including
+large-scale machine learning research.
 
 With its updated version of [Autograd](https://github.com/hips/autograd),
 JAX can automatically differentiate native
diff --git a/docs/index.rst b/docs/index.rst
index 6214d2bb4..bef957568 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1,8 +1,7 @@
 JAX: High-Performance Array Computing
 =====================================
 
-JAX is a Python library for accelerator-oriented array computation and program transformation,
-designed for high-performance numerical computing and large-scale machine learning.
+JAX is Autograd_ and XLA_, brought together for high-performance numerical computing.
 
 If you're looking to train neural networks, use Flax_ and start with its documentation.
 Some associated tools are Optax_ and Orbax_.
@@ -94,6 +93,8 @@ For an end-to-end transformer library built on JAX, see MaxText_.
    glossary
 
 
+.. _Autograd: https://github.com/hips/autograd
+.. _XLA: https://openxla.org/xla
 .. _Flax: https://flax.readthedocs.io/
 .. _Orbax: https://orbax.readthedocs.io/
 .. _Optax: https://optax.readthedocs.io/
diff --git a/jax/BUILD b/jax/BUILD
index 0ebfca2ba..8386660f3 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -1024,7 +1024,6 @@ pytype_strict_library(
     visibility = [":jax_extend_users"],
     deps = [
         "//jax/extend",
-        "//jax/extend:backend",
         "//jax/extend:core",
         "//jax/extend:linear_util",
         "//jax/extend:random",
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 9c112970a..522a3b779 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -772,7 +772,8 @@ def _flatten_bwd(in_tree, in_avals, out_trees, *args):
       results.append(Zero(ct.aval))
     else:
       if (not core.typecompat(a.at_least_vspace(), a_ := core.get_aval(ct))
-          and not _temporary_dtype_exception(a, a_)):
+          # TODO(mattjj): don't skip check with extended dtype tangent types
+          and not dtypes.issubdtype(a_.dtype, dtypes.extended)):
         msg = ("Custom VJP bwd rule must produce an output with the same "
                "shape/dtypes as the args tuple of the primal function, but at "
                f"output{keystr(kp)} the bwd rule produced an output of "
@@ -782,14 +783,6 @@ def _flatten_bwd(in_tree, in_avals, out_trees, *args):
       results.append(ct)
   yield results
 
-# TODO(mattjj): remove both these exceptions to cotangent compatibility check
-def _temporary_dtype_exception(a, a_) -> bool:
-  if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray):
-    return (a.shape == a_.shape and
-            (dtypes.issubdtype(a_.dtype, dtypes.extended) or
-             dtypes.issubdtype(a.dtype, dtypes.np.inexact)))
-  return False
-
 
 class CustomVJPCallPrimitive(core.CallPrimitive):
   initial_style: core.Primitive
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index b18af885d..cd0799297 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2076,8 +2076,8 @@ def lower_sharding_computation(
       any(not is_unspecified(o) for o in out_shardings))
 
   gs = GSPMDSharding.get_replicated(device_assignment)
-  if xla_extension_version < 241 or hasattr(backend, "compile_replicated"):
-    in_shardings = tuple(gs if is_unspecified(i) else i for i in in_shardings)
+  # if xla_extension_version < 241 or hasattr(backend, "compile_replicated"):
+  in_shardings = tuple(gs if is_unspecified(i) else i for i in in_shardings)
 
   da_object = _create_da_object(tuple(device_assignment))
 
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index b1e80b952..0e628fc5e 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -751,6 +751,8 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
   if dtypes.issubdtype(a.dtype, np.complexfloating):
     raise ValueError("quantile does not support complex input, as the operation is poorly defined.")
   if axis is None:
+    if keepdims:
+      keepdim = [1] * a.ndim
     a = a.ravel()
     axis = 0
   elif isinstance(axis, tuple):
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
index 984692f3f..2a0caaa95 100644
--- a/jax/_src/pallas/mosaic/lowering.py
+++ b/jax/_src/pallas/mosaic/lowering.py
@@ -55,6 +55,7 @@ from jax._src.util import safe_map
 from jax._src.util import safe_zip
 from jax._src.util import split_list
 from jax._src.util import unzip2
+from jax._src.util import unzip3
 from jax.experimental.mosaic.dialects import tpu
 import jax.numpy as jnp
 import numpy as np
@@ -745,71 +746,47 @@ def _maybe_cast_to_index(cast_to_index, x):
     return _make_index(x)
   return _ensure_mlir_value(x, aval=jax_core.ShapedArray((), jnp.int32))
 
-
-def _index_to_start_size_stride(
-    idx: tuple[indexing.Slice | int | ir.Value, ...], cast_to_index: bool
-) -> tuple[ir.Value, int, int, bool]:
+def _index_to_start_size(idx: tuple[indexing.Slice | int | ir.Value, ...],
+                         cast_to_index: bool) -> tuple[ir.Value, int, bool]:
   assert not isinstance(idx, slice)
   if isinstance(idx, indexing.Slice):
     start = _maybe_cast_to_index(cast_to_index, idx.start)
     size = idx.size
-    stride = idx.stride
     squeeze = False
   elif isinstance(idx, int):
     start = _maybe_cast_to_index(cast_to_index, idx)
     size = 1
-    stride = 1
     squeeze = True
   else:
     if np.shape(idx):
       raise ValueError(f"Can only use ()-shaped and slice indexing: {idx}")
     start = _maybe_cast_to_index(cast_to_index, idx)
     size = 1
-    stride = 1
     squeeze = True
-  return start, size, stride, squeeze
+  return start, size, squeeze
 
 
-def _indexer_to_start_size_stride(
-    indexer: NDIndexer,
-    ref_block_shape: tuple[int | pl_core.Mapped, ...],
-    *,
+def _indexer_to_start_size(
+    indexer: NDIndexer, ref_block_shape: tuple[int | pl_core.Mapped, ...], *,
     cast_to_index: bool,
-) -> tuple[
-    tuple[ir.Value, ...],
-    tuple[int, ...],
-    tuple[int, ...],
-    tuple[bool, ...],
-    tuple[int | pl_core.Mapped, ...],
-]:
+) -> tuple[tuple[ir.Value, ...], tuple[int, ...], tuple[bool, ...],
+           tuple[int | pl_core.Mapped, ...]]:
   indices_iter = iter(indexer.indices)
-  starts, sizes, strides, squeeze_dims = [], [], [], []
-  for s in ref_block_shape:
-    start, size, stride, squeeze_dim = (
-        (
-            _maybe_cast_to_index(cast_to_index, 0),
-            1,
-            1,
-            True,
-        )
-        if s is pl_core.mapped
-        else _index_to_start_size_stride(next(indices_iter), cast_to_index)
-    )
-    starts.append(start)
-    sizes.append(size)
-    strides.append(stride)
-    squeeze_dims.append(squeeze_dim)
+  starts, sizes, squeeze_dims = unzip3(
+      (
+          _maybe_cast_to_index(cast_to_index, 0),
+          1,
+          True,
+      )
+      if s is pl_core.mapped
+      else _index_to_start_size(next(indices_iter), cast_to_index)
+      for s in ref_block_shape
+  )
   next_index = next(indices_iter, None)
   assert next_index is None, (indexer.indices, ref_block_shape)
   new_ref_block_shape = tuple(s for s, squeeze in zip(sizes, squeeze_dims)
                               if not squeeze)
-  return (
-      tuple(starts),
-      tuple(sizes),
-      tuple(strides),
-      tuple(squeeze_dims),
-      new_ref_block_shape,
-  )
+  return tuple(starts), tuple(sizes), tuple(squeeze_dims), new_ref_block_shape
 
 
 def _slice_memref(ref: ir.Value, ref_aval: state.AbstractRef,
@@ -819,15 +796,9 @@ def _slice_memref(ref: ir.Value, ref_aval: state.AbstractRef,
                              tuple[int | pl_core.Mapped, ...]]:
   assert ref_block_shape is not None
   target_shape = indexer.get_indexer_shape()
-  starts, sizes, strides, squeeze_dims, ref_block_shape = (
-      _indexer_to_start_size_stride(
-          indexer,
-          ref_block_shape,
-          cast_to_index=False,
-      )
+  starts, sizes, squeeze_dims, ref_block_shape = _indexer_to_start_size(
+      indexer, ref_block_shape, cast_to_index=False,
   )
-  if not all((s is None or s == 1) for s in strides):
-    raise NotImplementedError("Strided slices of references are unsupported.")
   target_ref_ty = ir.MemRefType.get(
       tuple(sizes), _dtype_to_ir_type(ref_aval.dtype),
       memory_space=ref.type.memory_space)
@@ -875,21 +846,14 @@ def _load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree, **_):
       for a in idx_aval.indices
   ):
     raise ValueError("Cannot do int indexing on TPU")
-  starts, sizes, strides, _, _ = _indexer_to_start_size_stride(
-      idx,
-      ref_block_shape,
-      cast_to_index=True,
+  starts, sizes, _, _ = _indexer_to_start_size(
+      idx, ref_block_shape, cast_to_index=True,
   )
-  need_stride = not all((s is None or s == 1) for s in strides)
   load_aval = jax_core.ShapedArray(sizes, dtype=ref_aval.dtype)
   if is_smem_load:
     if ctx.avals_out[0].shape:
       raise ValueError("Can only load scalars from SMEM")
     return memref.LoadOp(ref, starts).result
-  if need_stride:
-    load_val = tpu.StridedLoadOp(
-        aval_to_ir_type(load_aval), ref, starts, strides
-    ).result
   else:
     load_val = vector.LoadOp(aval_to_ir_type(load_aval), ref, starts).result
   if load_aval == aval_out:
@@ -932,12 +896,10 @@ def _masked_swap_lowering_rule(
     raise NotImplementedError(
         "Indexing into a ()-shaped Ref not yet supported on TPU.")
 
-  starts, _, strides, _, _ = _indexer_to_start_size_stride(
-      idx,
-      ref_block_shape,
-      cast_to_index=True,
+  starts, _, _, _ = _indexer_to_start_size(
+      idx, ref_block_shape, cast_to_index=True,
   )
-  need_stride = not all((s is None or s == 1) for s in strides)
+
   if is_smem_store:
     if val_aval.shape:
       raise ValueError("Can only store scalars to SMEM")
@@ -956,10 +918,7 @@ def _masked_swap_lowering_rule(
   mem_aval = aval_out.update(shape=tuple(mem_slice_shape))
   mem_aval_vec_type = ir.VectorType.get(mem_aval.shape,
                                         _dtype_to_ir_type(mem_aval.dtype))
-  if need_stride:
-    result = tpu.StridedLoadOp(mem_aval_vec_type, ref, starts, strides).result
-  else:
-    result = vector.LoadOp(mem_aval_vec_type, ref, starts).result
+  result = vector.LoadOp(mem_aval_vec_type, ref, starts).result
   if mem_aval != aval_out:
     # We are slicing a scalar so provided dummy 1 indices
     result_vec_type = ir.VectorType.get(aval_out.shape,
@@ -968,10 +927,7 @@ def _masked_swap_lowering_rule(
     val_vec_type = ir.VectorType.get(mem_aval.shape,
                                      _dtype_to_ir_type(mem_aval.dtype))
     val = vector.ShapeCastOp(val_vec_type, val).result
-  if need_stride:
-    tpu.StridedStoreOp(val, ref, starts, strides)
-  else:
-    vector.StoreOp(val, ref, starts)
+  vector.StoreOp(val, ref, starts)
   return result
 
 
diff --git a/jax/_src/state/indexing.py b/jax/_src/state/indexing.py
index 60b99f100..175d88398 100644
--- a/jax/_src/state/indexing.py
+++ b/jax/_src/state/indexing.py
@@ -33,19 +33,16 @@ class Slice:
   """Represents a slice with a dynamic start index and a fixed size."""
   start: Any
   size: int
-  stride: int = 1
 
   def __post_init__(self):
     if self.size < 0:
       raise ValueError("`size` must not be negative.")
-    if self.stride < 1:
-      raise ValueError("`stride` must be >= 1.")
 
   def tree_flatten(self):
     # If `start` is statically known, we treat it as static information
     if isinstance(self.start, int):
-      return (), (self.start, self.size, self.stride)
-    return (self.start,), (self.size, self.stride)
+      return (), (self.start, self.size)
+    return (self.start,), (self.size,)
 
   @classmethod
   def tree_unflatten(cls, aux_data, children) -> Slice:
@@ -54,30 +51,21 @@ class Slice:
   @classmethod
   def from_slice(cls, slc: slice, size: int) -> Slice:
     start, stop, step = slc.indices(size)
-    if step < 1:
-      raise ValueError(f"slice must have a step >= 1 (found: {step})")
-    return cls(start, max((stop - start + step - 1) // step, 0), step)
+    if step != 1:
+      raise ValueError(f"slice must have a step of 1 (found: {step})")
+    return cls(start, max(stop - start, 0))
 
 
-def dslice(
-    start: int | Array | None,
-    size: int | None = None,
-    stride: int | None = None,
-) -> slice | Slice:
+def dslice(start: int |  Array | None, size: int | None = None
+           ) -> slice | Slice:
   """Constructs a `Slice` from a start and a size."""
   if start is None:
     return slice(None)
-  if stride is None:
-    stride = 1
-  if not isinstance(stride, int):
-    raise ValueError("Non-static stride in `dslice`")
   if size is None:
     if not isinstance(start, int):
       raise ValueError("Non-static `dslice`")
-    return Slice(0, start, stride)
-  return Slice(start, size, stride)
-
-
+    return Slice(0, start)
+  return Slice(start, size)
 ds = dslice  # Handy alias
 
 
@@ -125,10 +113,9 @@ class NDIndexer:
         if value := _maybe_concretize(start):
           if value >= s:
             raise ValueError(f"Out of bound slice: start={value}, dim={s}.")
-          if value + (idx.size - 1) * idx.stride >= s:
+          if value + idx.size > s:
             raise ValueError(
-                f"Out of bound slice: start={value}, size={idx.size},"
-                f" stride={idx.stride}, dim={s}."
+                f"Out of bound slice: start={value}, size={idx.size}, dim={s}."
             )
         continue
       # The shape of indexer integers should be broadcastable up to the
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index f33aa8938..71250a34c 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -40,7 +40,6 @@ from jax._src import linear_util as lu
 from jax._src import ops
 from jax._src import pjit
 from jax._src import prng
-from jax._src import random
 from jax._src import sharding_impls
 from jax._src import source_info_util
 from jax._src import traceback_util
@@ -942,8 +941,7 @@ for o in it.chain(lax.__dict__.values(), slicing.__dict__.values(),
                   special.__dict__.values(), convolution.__dict__.values(),
                   fft.__dict__.values(), linalg.__dict__.values(),
                   ops.__dict__.values(), ad_util.__dict__.values(),
-                  prng.__dict__.values(), ann.__dict__.values(),
-                  random.__dict__.values()):
+                  prng.__dict__.values(), ann.__dict__.values()):
   if isinstance(o, core.Primitive):
     register_standard_check(o)
     register_standard_rewrite(o)
diff --git a/jax/extend/BUILD b/jax/extend/BUILD
index 9b24200e0..7b367d3a2 100644
--- a/jax/extend/BUILD
+++ b/jax/extend/BUILD
@@ -26,7 +26,6 @@ pytype_strict_library(
     name = "extend",
     srcs = ["__init__.py"],
     deps = [
-        ":backend",
         ":core",
         ":linear_util",
         ":random",
@@ -46,12 +45,6 @@ pytype_strict_library(
     deps = ["//jax:core"],
 )
 
-pytype_strict_library(
-    name = "backend",
-    srcs = ["backend.py"],
-    deps = ["//jax"],
-)
-
 pytype_strict_library(
     name = "random",
     srcs = ["random.py"],
diff --git a/jax/extend/__init__.py b/jax/extend/__init__.py
index 3f4327dde..77c81488c 100644
--- a/jax/extend/__init__.py
+++ b/jax/extend/__init__.py
@@ -29,7 +29,6 @@ Breaking changes will be announced via the
 """
 
 from jax.extend import (
-    backend as backend,
     core as core,
     linear_util as linear_util,
     random as random,
diff --git a/jax/extend/backend.py b/jax/extend/backend.py
deleted file mode 100644
index 7aa2c8a06..000000000
--- a/jax/extend/backend.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# Copyright 2024 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Note: import <name> as <name> is required for names to be exported.
-# See PEP 484 & https://github.com/google/jax/issues/7570
-
-from jax._src.api import (
-  clear_backends as clear_backends,
-)
diff --git a/jaxlib/mosaic/dialect/tpu/tpu.td b/jaxlib/mosaic/dialect/tpu/tpu.td
index c0d1db245..68d24ed8f 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu.td
+++ b/jaxlib/mosaic/dialect/tpu/tpu.td
@@ -197,12 +197,12 @@ def TPU_LoadOp : TPU_Op<"load"> {
 def TPU_StridedLoadOp : TPU_Op<"strided_load"> {
   let arguments = (ins
     AnyMemRef:$base,
-    Variadic<Index>:$indices,
+    DenseI32ArrayAttr:$indices,
     DenseI32ArrayAttr:$strides
   );
   let results = (outs AnyVector:$result);
   let assemblyFormat = [{
-    $base `[` $indices `]` attr-dict `:` type($base) `,` type($result)
+    $base attr-dict `:` type($base) `,` type($result)
   }];
   let hasVerifier = 1;
 }
@@ -211,12 +211,12 @@ def TPU_StridedStoreOp : TPU_Op<"strided_store"> {
   let arguments = (ins
     AnyVector:$valueToStore,
     AnyMemRef:$base,
-    Variadic<Index>:$indices,
+    DenseI32ArrayAttr:$indices,
     DenseI32ArrayAttr:$strides
   );
   let results = (outs);
   let assemblyFormat = [{
-    $base `[` $indices `]` `,` $valueToStore attr-dict `:` type($base) `,` type($valueToStore)
+    $base `,` $valueToStore attr-dict `:` type($base) `,` type($valueToStore)
   }];
   let hasVerifier = 1;
 }
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_ops.cc b/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
index fd8785684..53941cbd2 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
+++ b/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
@@ -202,10 +202,21 @@ LogicalResult verifyStridedOp(Op op, MemRefType memref_ty,
     return failure();
   }
   for (int64_t i = 0; i < memref_ty.getRank(); ++i) {
+    if (indices[i] < 0 && indices[i] >= memref_ty.getDimSize(i)) {
+      op.emitError("Indices[")
+          << i << "]=" << indices[i] << " is out of range [0, "
+          << memref_ty.getDimSize(i) << ")";
+      return failure();
+    }
     if (strides[i] < 1) {
       op.emitError("Strides[") << i << "]=" << strides[i] << " must be >= 1";
       return failure();
     }
+    if ((indices[i] + (vector_ty.getDimSize(i) - 1) * strides[i]) >
+        memref_ty.getDimSize(i)) {
+      op.emitError() << "Strided slice is out of range at dim " << i;
+      return failure();
+    }
   }
   return success();
 }
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
index e02fdb365..e878ebd58 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -1159,9 +1159,9 @@ LogicalResult tpu_load_rule(RewriteContext &ctx, Operation &op,
 }
 
 LogicalResult strided_op_rule_impl(RewriteContext &ctx, Operation &op,
-                                   Value base_ref, ValueRange indices,
-                                   const VectorType &vty,
+                                   Value base_ref, const VectorType &vty,
                                    const VectorLayout &layout,
+                                   const ArrayRef<int32_t> &indices,
                                    const ArrayRef<int32_t> &strides) {
   if (!isa<tpu::StridedLoadOp, tpu::StridedStoreOp>(op)) {
     return op.emitOpError("Not implemented: Unsupported strided op")
@@ -1198,10 +1198,7 @@ LogicalResult strided_op_rule_impl(RewriteContext &ctx, Operation &op,
   if (strides[rank - 1] != 1) {
     return op.emitOpError("Not Implemented: Stride on last dim is not 1");
   }
-  auto last_idx = getIntConst(indices[rank - 1], /*silent=*/true);
-  if (failed(last_idx)) {
-    return op.emitOpError("Not Implemented: Dynamic index on last dim");
-  } else if (last_idx.value() != 0) {
+  if (indices[rank - 1] != 0) {
     return op.emitOpError("Not Implemented: Index on last dim is not 0");
   }
   ImplicitLocOpBuilder builder(op.getLoc(), &op);
@@ -1227,8 +1224,8 @@ LogicalResult strided_op_rule_impl(RewriteContext &ctx, Operation &op,
       int64_t stride = (i < rank - 2)
                            ? strides[i]
                            : (strides[i] * ctx.target_shape[i - rank + 2]);
-      idxs[i] = builder.create<arith::AddIOp>(
-          indices[i], IdxConst(tile_idxs[i] * stride, builder, op.getLoc()));
+      idxs[i] =
+          IdxConst(indices[i] + tile_idxs[i] * stride, builder, op.getLoc());
     }
     SmallVector<bool> sublane_mask(ctx.target_shape[0], true);
     int64_t sublane_rem = vty.getDimSize(rank - 2) % ctx.target_shape[0];
@@ -1267,9 +1264,12 @@ LogicalResult tpu_strided_load_rule(RewriteContext &ctx, Operation &op,
   TPU_ASSERT_OP(layouts_out.front().has_value());
   const VectorLayout &layout_out = *layouts_out.front();
   auto load_op = cast<tpu::StridedLoadOp>(op);
+  const auto base_ref = load_op.getBase();
+  const auto indices = load_op.getIndices();
+  const auto strides = load_op.getStrides();
   const auto vty = cast<VectorType>(load_op.getResult().getType());
-  return strided_op_rule_impl(ctx, op, load_op.getBase(), load_op.getIndices(),
-                              vty, layout_out, load_op.getStrides());
+  return strided_op_rule_impl(ctx, op, base_ref, vty, layout_out, indices,
+                              strides);
 }
 
 // TODO(jevinjiang): maybe unify with vector store?
@@ -1283,10 +1283,12 @@ LogicalResult tpu_strided_store_rule(RewriteContext &ctx, Operation &op,
 
   const VectorLayout &to_store_layout = *layouts_in.front();
   auto store_op = cast<tpu::StridedStoreOp>(op);
+  const auto base_ref = store_op.getBase();
+  const auto indices = store_op.getIndices();
+  const auto strides = store_op.getStrides();
   const auto vty = store_op.getValueToStore().getType();
-  return strided_op_rule_impl(ctx, op, store_op.getBase(),
-                              store_op.getIndices(), vty, to_store_layout,
-                              store_op.getStrides());
+  return strided_op_rule_impl(ctx, op, base_ref, vty, to_store_layout, indices,
+                              strides);
 }
 
 LogicalResult matmul_rule_impl(RewriteContext &ctx, Operation &op,
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/debug_assert_insertion.cc b/jaxlib/mosaic/dialect/tpu/transforms/debug_assert_insertion.cc
index 8e1246299..1eb7c12a8 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/debug_assert_insertion.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/debug_assert_insertion.cc
@@ -55,11 +55,9 @@ rule_type as_generic_rule(void (*rule)(Op)) {
 
 void assertIsValidSubwindow(Operation *op, mlir::ValueRange base_indices,
                             ArrayRef<int64_t> window_shape,
-                            ArrayRef<int64_t> full_shape,
-                            ArrayRef<int32_t> strides = {}) {
+                            ArrayRef<int64_t> full_shape) {
   if (base_indices.size() != window_shape.size() ||
-      base_indices.size() != full_shape.size() ||
-      (!strides.empty() && base_indices.size() != strides.size())) {
+      base_indices.size() != full_shape.size()) {
     return;  // Malformed op.
   }
   if (base_indices.empty()) {
@@ -70,15 +68,14 @@ void assertIsValidSubwindow(Operation *op, mlir::ValueRange base_indices,
   for (auto [dim, access] :
        llvm::enumerate(llvm::zip(base_indices, window_shape, full_shape))) {
     auto [idx, size, bound] = access;
-    int64_t stride = strides.empty() ? 1 : strides[dim];
     Value positive = builder.create<arith::CmpIOp>(
         arith::CmpIPredicate::sge, idx,
         builder.create<arith::ConstantOp>(builder.getIntegerAttr(idx_type, 0)));
     Value in_bounds = builder.create<arith::CmpIOp>(
-        arith::CmpIPredicate::slt,
+        arith::CmpIPredicate::sle,
         builder.create<arith::AddIOp>(
             idx, builder.create<arith::ConstantOp>(
-                     builder.getIntegerAttr(idx_type, (size - 1) * stride))),
+                     builder.getIntegerAttr(idx_type, size))),
         builder.create<arith::ConstantOp>(
             builder.getIntegerAttr(idx_type, bound)));
     std::string msg;
@@ -110,21 +107,6 @@ void tpu_memref_slice_rule(tpu::MemRefSliceOp op) {
                          /*full_shape=*/op.getMemRef().getType().getShape());
 }
 
-void tpu_strided_load_rule(tpu::StridedLoadOp op) {
-  assertIsValidSubwindow(op, op.getIndices(),
-                         /*window_shape=*/op.getResult().getType().getShape(),
-                         /*full_shape=*/op.getBase().getType().getShape(),
-                         /*strides=*/op.getStrides());
-}
-
-void tpu_strided_store_rule(tpu::StridedStoreOp op) {
-  assertIsValidSubwindow(
-      op, op.getIndices(),
-      /*window_shape=*/op.getValueToStore().getType().getShape(),
-      /*full_shape=*/op.getBase().getType().getShape(),
-      /*strides=*/op.getStrides());
-}
-
 const llvm::StringMap<rule_type> &rules() {
   static auto rules = new llvm::StringMap<rule_type>{
       // TODO: tpu::LoadOp, tpu::StoreOp
@@ -132,10 +114,6 @@ const llvm::StringMap<rule_type> &rules() {
       {vector::StoreOp::getOperationName(), as_generic_rule(vector_store_rule)},
       {tpu::MemRefSliceOp::getOperationName(),
        as_generic_rule(tpu_memref_slice_rule)},
-      {tpu::StridedLoadOp::getOperationName(),
-       as_generic_rule(tpu_strided_load_rule)},
-      {tpu::StridedStoreOp::getOperationName(),
-       as_generic_rule(tpu_strided_store_rule)},
   };
   return *rules;
 }
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
index a69bcc3d8..59f1a9965 100644
--- a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
+++ b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -616,9 +616,7 @@ class VectorLayoutInferer {
     }
     auto store_layout = VectorLayout(bitwidth, {0, 0}, nativeTiling(bitwidth),
                                      ImplicitDim::kNone);
-    SmallVector<Layout, 5> in_layout{op->getNumOperands(), kNoLayout};
-    in_layout[0] = store_layout;
-    setInLayout(op, in_layout);
+    setInLayout(op, {store_layout, kNoLayout});
     return success();
   }
 


Test patch:
diff --git a/tests/extend_test.py b/tests/extend_test.py
index b49c1ac09..37f5c911d 100644
--- a/tests/extend_test.py
+++ b/tests/extend_test.py
@@ -18,7 +18,6 @@ import jax
 import jax.extend as jex
 import jax.numpy as jnp
 
-from jax._src import api
 from jax._src import abstract_arrays
 from jax._src import linear_util
 from jax._src import prng
@@ -40,7 +39,6 @@ class ExtendTest(jtu.JaxTestCase):
     self.assertIs(jex.random.unsafe_rbg_prng_impl, prng.unsafe_rbg_prng_impl)
 
     # Assume these are tested elsewhere, only check equivalence
-    self.assertIs(jex.backend.clear_backends, api.clear_backends)
     self.assertIs(jex.core.array_types, abstract_arrays.array_types)
     self.assertIs(jex.linear_util.StoreException, linear_util.StoreException)
     self.assertIs(jex.linear_util.WrappedFun, linear_util.WrappedFun)
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index 0b8bbd79e..b72e47d9d 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -662,6 +662,7 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     [dict(a_shape=a_shape, axis=axis)
       for a_shape, axis in (
         ((7,), None),
+        ((6, 7,), None),
         ((47, 7), 0),
         ((47, 7), ()),
         ((4, 101), 1),
@@ -716,6 +717,7 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
     [dict(a_shape=a_shape, axis=axis)
       for a_shape, axis in (
         ((7,), None),
+        ((6, 7,), None),
         ((47, 7), 0),
         ((4, 101), 1),
       )
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index eccbd1289..a3f17e7b6 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -3942,22 +3942,6 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertArraysEqual(out4, np_inp * 3)
     self.assertArraysEqual(out5, np_inp.T)
 
-  def test_parameter_tupled_jit(self):
-    if not jtu.test_device_matches(["tpu"]):
-      self.skipTest('Parameters are tupled only on TPU if >2000 parameters')
-
-    mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
-    s = NamedSharding(mesh, P('x'))
-
-    @jax.jit
-    def f(*args):
-      return args * 2
-
-    inp = np.arange(8)
-    arr = jax.device_put(inp, s)
-    inps = [arr, *[inp] * 2001]
-    f(inps)  # doesn't crash
-
 
 class TempSharding(Sharding):
 
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index d84637723..b25e33ca2 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -30,7 +30,6 @@ from jax._src import maps
 from jax._src import test_util as jtu
 from jax._src import util
 from jax._src.lib import xla_client
-from jax._src.lib import xla_extension_version
 from jax.experimental import io_callback
 from jax.experimental import pjit
 from jax.experimental.maps import xmap
@@ -640,6 +639,7 @@ class PureCallbackTest(jtu.JaxTestCase):
     out = h(jnp.arange(4.)[None], 4.)
     np.testing.assert_allclose(out, np.sin(np.arange(4.)[None]) + 4.)
 
+
   def test_vmap_vectorized_callback_errors_if_returns_wrong_shape(self):
 
     def cb(x):
@@ -716,18 +716,6 @@ class PureCallbackTest(jtu.JaxTestCase):
         ValueError, "Pure callbacks do not support JVP."):
       f(2.)
 
-  @unittest.skipIf(xla_extension_version < 245, "jaxlib version too old")
-  def test_error_propagation(self):
-    def throws_error_fn(x):
-      raise RuntimeError("Errors should propagate.")
-
-    @jax.jit
-    def f(x):
-      return jax.pure_callback(throws_error_fn, x, x)
-
-    with self.assertRaisesRegex(Exception, "Errors should propagate."):
-      print(np.array(f(2.0)), flush=True)
-
   def test_can_take_grad_of_pure_callback_with_custom_jvp(self):
 
     @jax.custom_jvp
@@ -845,6 +833,7 @@ class PureCallbackTest(jtu.JaxTestCase):
     # callback alive.
     np.testing.assert_allclose(out, np.full((num_devices, 4), 11, np.float32))
 
+
   def test_callback_inside_xmap(self):
 
     def _callback(x):