Repository: ml-explore/mlx (2026-05-29T15:28:48Z)

Issue:
Signed int-overflow (UB) in shape-op size/rank arithmetic (roll, tile, flatten) for large/degenerate shapes
### Summary

Several shape ops compute sizes / ranks / products in 32-bit `int` without overflow checks, so large-or-degenerate (but otherwise valid, non-negative, `int32`-fitting) shapes trigger **C++ signed-integer overflow (UB)** before any error is returned. These are reachable from the safe high-level API.

### Affected paths (function names; line refs from the version vendored via `mlx-c`, ~v0.31.x)

- **`roll`** — `array roll(const array&, const Shape& shift, ...)` folds shifts into an unchecked `int total_shift` (`mlx/ops.cpp`), and the per-axis form computes `(sh < 0) ? (-sh) % size : ...` which negates `INT_MIN` (UB) for a shift of `i32::MIN`.
- **`tile`** — builds output dims via `reps[i] * shape[i]` in `int` (`mlx/ops.cpp`), and constructs intermediate `expand`/`broadcast` shapes of rank `aligned_rank + count(reps != 1)`, which can exceed `INT_MAX` and then index `shape.size()` with `int`.
- **`flatten` / `Flatten::output_shape`** — multiplies dimensions left-to-right in an `int` accumulator (`mlx/primitives.cpp`); a degenerate shape like `[INT_MAX, 2, 0]` overflows on `INT_MAX * 2` before reaching the `0`. (Reached transitively, e.g. by the no-axis `roll`.)
- `broadcast_shapes` (`mlx/utils.cpp`) narrows `size_t` rank to `int`; `ArrayDesc::init` (`mlx/array.cpp`) loops `for (int i = shape.size()-1; ...)`.

### Minimal repros (Python)

```python
import mlx.core as mx

# flatten product overflow (cheap: 0-element lazy array, no large alloc)
mx.eval(mx.roll(mx.zeros([2**31 - 1, 2, 0]), 1))

# tile reps*dim overflow
mx.eval(mx.tile(mx.zeros([2]), [2**31 - 1]))

# INT_MIN shift negation
mx.eval(mx.roll(mx.zeros([4]), -2**31))
```

Building MLX with `-fsanitize=undefined` should flag the signed-overflow in these paths.

### Impact

Undefined behavior reachable from the public API on valid (non-negative, `int32`-range) shape/argument values — not just from already-invalid inputs. Bindings (Swift / the Rust `mlxrs` wrapper this was found in) that forward safe inputs inherit the UB and cannot soundly prevent it without re-implementing MLX's internal shape arithmetic.

### Suggested fix

Use 64-bit / checked arithmetic for shape products, ranks, and shift sums in the shape-op size computations (and reject `INT_MIN` negation), returning a graceful error instead of overflowing.


Discussion:

Patch:
diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 6ad41e2e38..349a3736c5 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -1349,7 +1349,8 @@ array tile(
     }
     expand_shape.push_back(shape[i]);
     broad_shape.push_back(shape[i]);
-    final_shape.push_back(reps[i] * shape[i]);
+    final_shape.push_back(
+        safe_cast(static_cast<int64_t>(reps[i]) * shape[i], "tile"));
   }
 
   auto x = reshape(arr, std::move(expand_shape), s);
@@ -6350,7 +6351,10 @@ array roll(
     if (size == 0) {
       continue; // skip rolling this axis if it has size 0
     }
-    auto split_index = (sh < 0) ? (-sh) % size : size - sh % size;
+    // Promote to 64-bit so negating a shift of INT_MIN does not overflow.
+    int64_t sh64 = sh;
+    auto split_index = static_cast<ShapeElem>(
+        (sh64 < 0) ? (-sh64) % size : size - sh64 % size);
 
     auto parts = split(result, Shape{split_index}, ax, s);
     std::swap(parts[0], parts[1]);
@@ -6369,11 +6373,11 @@ array roll(const array& a, int shift, StreamOrDevice s /* = {} */) {
 }
 
 array roll(const array& a, const Shape& shift, StreamOrDevice s /* = {} */) {
-  int total_shift = 0;
-  for (auto& s : shift) {
-    total_shift += s;
+  int64_t total_shift = 0;
+  for (auto& sh : shift) {
+    total_shift += sh;
   }
-  return roll(a, total_shift, s);
+  return roll(a, safe_cast(total_shift, "roll"), s);
 }
 
 array roll(const array& a, int shift, int axis, StreamOrDevice s /* = {} */) {
@@ -6394,11 +6398,12 @@ array roll(
     const Shape& shift,
     int axis,
     StreamOrDevice s /* = {} */) {
-  int total_shift = 0;
-  for (auto& s : shift) {
-    total_shift += s;
+  int64_t total_shift = 0;
+  for (auto& sh : shift) {
+    total_shift += sh;
   }
-  return roll(a, Shape{total_shift}, std::vector<int>{axis}, s);
+  return roll(
+      a, Shape{safe_cast(total_shift, "roll")}, std::vector<int>{axis}, s);
 }
 
 array real(const array& a, StreamOrDevice s /* = {} */) {


Test patch:
diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp
index 500195093e..c2230f77cf 100644
--- a/tests/ops_tests.cpp
+++ b/tests/ops_tests.cpp
@@ -4284,3 +4284,24 @@ TEST_CASE("test max min with nan") {
   CHECK(array_equal(max_result, expected, true).item<bool>());
   CHECK(array_equal(min_result, expected, true).item<bool>());
 }
+
+TEST_CASE("roll and tile shape overflow") {
+  // Shape arithmetic must not overflow (signed-int UB) for large but otherwise
+  // valid int32 inputs; out-of-range results are rejected gracefully.
+  // https://github.com/ml-explore/mlx/issues/3601
+
+  // tile: reps * dim exceeding int32 raises instead of overflowing.
+  CHECK_THROWS_AS(tile(zeros({2}), {2147483647}), std::overflow_error);
+
+  // roll: a shift sum exceeding int32 raises instead of overflowing.
+  CHECK_THROWS_AS(
+      roll(zeros({4}), Shape{2147483647, 2147483647}), std::overflow_error);
+  CHECK_THROWS_AS(
+      roll(zeros({4}), Shape{2147483647, 2147483647}, 0), std::overflow_error);
+
+  // roll: a shift of INT_MIN must not negate-overflow. INT_MIN mod 4 == 0, so
+  // rolling a size-4 axis by INT_MIN is the identity.
+  auto x = array({1, 2, 3, 4});
+  auto rolled = roll(x, -2147483647 - 1);
+  CHECK(array_equal(rolled, x).item<bool>());
+}