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

[uqio]

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)

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.