[Mosaic] Fix bug in divisibility check in infer_vector_layout load an…

…d store rules PiperOrigin-RevId: 609876232
jax-ml · Feb 24, 2024 · 61aa7e8 · 61aa7e8
1 parent 072b43b
commit 61aa7e8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc
@@ -847,7 +847,7 @@ class VectorLayoutInferer {
         tile_offsets.push_back(cast<IntegerAttr>(cst_op.getValue()).getInt() %
                                tiling[i]);
       } else {
-        if (failed(verifyDivisibleIndex(tiled_index, tiling[0], dim, op))) {
+        if (failed(verifyDivisibleIndex(tiled_index, tiling[i], dim, op))) {
           return failure();
         }
         tile_offsets.push_back(0);
@@ -1188,7 +1188,7 @@ class VectorLayoutInferer {
         tile_offsets.push_back(cast<IntegerAttr>(cst_op.getValue()).getInt() %
                                tiling[i]);
       } else {
-        if (failed(verifyDivisibleIndex(tiled_index, tiling[0], dim, op))) {
+        if (failed(verifyDivisibleIndex(tiled_index, tiling[i], dim, op))) {
           return failure();
         }
         tile_offsets.push_back(0);