[Mosaic TPU] Allow stride 0 in strided_load for broadcast.

PiperOrigin-RevId: 835372731

Author: megrez-yliu

jax/_src/state/indexing.py

@@ -42,8 +42,8 @@ class Slice:
   stride: int = 1
 
   def __post_init__(self):
-    if self.stride < 1:
-      raise ValueError("`stride` must be >= 1.")
+    if self.stride < 0:
+      raise ValueError("`stride` must be >= 0.")
 
   @property
   def is_dynamic_start(self):

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -735,7 +735,7 @@ LogicalResult MemRefBitcastOp::canonicalize(MemRefBitcastOp op,
 
 template <typename Op>
 LogicalResult verifyStridedOp(Op op, MemRefType memref_ty,
-                              VectorType vector_ty) {
+                              VectorType vector_ty, int64_t min_stride) {
   auto indices = op.getIndices();
   auto strides = op.getStrides();
   if (memref_ty.getRank() != indices.size()) {
@@ -754,8 +754,9 @@ LogicalResult verifyStridedOp(Op op, MemRefType memref_ty,
     return failure();
   }
   for (int64_t i = 0; i < memref_ty.getRank(); ++i) {
-    if (strides[i] < 1) {
-      op.emitError("Strides[") << i << "]=" << strides[i] << " must be >= 1";
+    if (strides[i] < min_stride) {
+      op.emitError("Strides[") << i << "]=" << strides[i] << " must be >= "
+          << min_stride;
       return failure();
     }
   }
@@ -764,12 +765,13 @@ LogicalResult verifyStridedOp(Op op, MemRefType memref_ty,
 
 LogicalResult StridedLoadOp::verify() {
   return verifyStridedOp<StridedLoadOp>(*this, getMemRefType(getBase()),
-                                        getType());
+                                        getType(), /*min_stride=*/0);
 }
 
 LogicalResult StridedStoreOp::verify() {
   return verifyStridedOp<StridedStoreOp>(*this, getMemRefType(getBase()),
-                                         getValueToStore().getType());
+                                         getValueToStore().getType(),
+                                         /*min_stride=*/1);
 }
 
 template <typename Op>

tests/pallas/indexing_test.py

@@ -541,6 +541,46 @@ def body(x_ref, y_ref1, y_ref2):
         y2[slices], expected, err_msg="Strided Store Error"
     )
 
+  @hp.given(hps.data())
+  def test_load_and_broadcast_with_stride_0(self, data):
+    if not jtu.if_cloud_tpu_at_least(2025, 11, 25):
+      self.skipTest("Requires libtpu built after 2025-11-25")
+    if self.INTERPRET:
+      self.skipTest("TODO: fails in interpret mode.")
+    dtype = jnp.float32
+    rank = data.draw(hps.integers(min_value=2, max_value=4))
+    shape = data.draw(hps.tuples(
+        *(hps.integers(min_value=1, max_value=10) for _ in range(rank - 1))))
+    shape = (*shape, 128)
+
+    strides = data.draw(hps.tuples(
+        *(hps.sampled_from([0, 1]) for _ in range(rank - 1))))
+    strides = (*strides, 1)
+
+    indices = []
+    for i in range(rank):
+      index = (data.draw(hps.integers(min_value=0, max_value=shape[i] - 1))
+               if strides[i] == 0 else 0)
+      indices.append(index)
+
+    def body(x_ref, y_ref):
+      slices = tuple(
+          pl.ds(i, l, s) for i, l, s in zip(indices, shape, strides)
+      )
+      y_ref[...] = x_ref[slices]
+
+    x = random.normal(random.key(33), shape, dtype=dtype)
+    y = self.pallas_call(
+        body,
+        out_shape=jax.ShapeDtypeStruct(shape, dtype),
+    )(x)
+    slices = tuple(slice(i, l, 1) if s != 0 else slice(i, i + 1, 1)
+                   for i, l, s in zip(indices, shape, strides))
+
+    expected = jnp.broadcast_to(x[slices], shape)
+    self.assertAllClose(y, expected)
+
+
   def test_load_with_dynamic_2nd_minor_index(self):
     if pltpu is None:
       self.skipTest("No TPU module available.")