[Mosaic][util] Fix ComputeTileStrides for 0D tiles

Previously, `tiling.size() - 1` could underflow for 0D tiles, leading to and OOB access and crash

PiperOrigin-RevId: 833506958

Author: tlongeri

jaxlib/mosaic/dialect/tpu/util.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <array>
+#include <cstddef>
 #include <cstdint>
 #include <optional>
 #include <ostream>
@@ -53,13 +54,14 @@ std::ostream &operator<<(std::ostream &os, Print p) {
 
 SmallVector<int64_t> ComputeTileStrides(absl::Span<const int64_t> shape,
                                         absl::Span<const int64_t> tiling) {
+  CHECK_LE(tiling.size(), shape.size());
   SmallVector<int64_t> tile_strides(shape.size());
   int64_t stride = 1;
-  for (int64_t i = 0; i < shape.size(); ++i) {
-    int64_t idx = shape.size() - 1 - i;
-    int64_t tiling_idx = tiling.size() - 1 - i;
+  for (size_t i = 0; i < shape.size(); ++i) {
+    const size_t idx = shape.size() - 1 - i;
     tile_strides[idx] = stride;
-    if (tiling_idx >= 0) {
+    if (i < tiling.size()) {
+      const size_t tiling_idx = tiling.size() - 1 - i;
       stride *= llvm::divideCeil(shape[idx], tiling[tiling_idx]);
     } else {
       stride *= shape[idx];