Fix int32 overflow in concatenate/repeat/kron

mlx/ops.cpp

@@ -1330,7 +1330,9 @@ array concatenate(
   };
 
   auto shape = arrays[0].shape();
-  shape[ax] = 0;
+  // Accumulate the concatenation axis in 64 bits so a total that does not fit
+  // in a shape dimension is reported rather than silently wrapping.
+  int64_t concat_size = 0;
   // Make the output shape and validate that all arrays have the same shape
   // except for the concatenation axis.
   for (auto& a : arrays) {
@@ -1349,8 +1351,9 @@ array concatenate(
         throw_invalid_shapes();
       }
     }
-    shape[ax] += a.shape(ax);
+    concat_size += a.shape(ax);
   }
+  shape[ax] = safe_cast(concat_size, "concatenate");
 
   // Promote all the arrays to the same type
   auto dtype = result_type(arrays);
@@ -1426,7 +1429,8 @@ array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) {
 
   // Reshape back into a contiguous array where S_axis is now S_axis * repeats
   shape.erase(shape.begin() + axis + 1);
-  shape[axis] *= repeats;
+  shape[axis] =
+      safe_cast(static_cast<int64_t>(shape[axis]) * repeats, "repeat");
   out = reshape(out, shape, s);
 
   return out;
@@ -3684,11 +3688,13 @@ array kron(const array& a, const array& b, StreamOrDevice s /* = {} */) {
 
   for (int i = ndim - 1, j = a.ndim() - 1; j >= 0; j--, i--) {
     a_shape[2 * i] = a.shape(j);
-    out_shape[i] *= a.shape(j);
+    out_shape[i] =
+        safe_cast(static_cast<int64_t>(out_shape[i]) * a.shape(j), "kron");
   }
   for (int i = ndim - 1, j = b.ndim() - 1; j >= 0; j--, i--) {
     b_shape[2 * i + 1] = b.shape(j);
-    out_shape[i] *= b.shape(j);
+    out_shape[i] =
+        safe_cast(static_cast<int64_t>(out_shape[i]) * b.shape(j), "kron");
   }
 
   return reshape(

mlx/primitives.cpp

@@ -1187,9 +1187,11 @@ bool Concatenate::is_equivalent(const Primitive& other) const {
 std::vector<Shape> Concatenate::output_shapes(
     const std::vector<array>& inputs) {
   auto shape = inputs[0].shape();
+  int64_t concat_size = shape[axis_];
   for (int i = 1; i < inputs.size(); ++i) {
-    shape[axis_] += inputs[i].shape(axis_);
+    concat_size += inputs[i].shape(axis_);
   }
+  shape[axis_] = safe_cast(concat_size, "concatenate");
   return {std::move(shape)};
 }
 

python/tests/test_ops.py

@@ -126,6 +126,27 @@ def test_shape_overflow_error(self):
             mx.broadcast_to(a, [too_big, 1])
         self.assertIn(str(too_big), str(cm.exception))
 
+        # A concatenation axis that does not fit is computed rather than given,
+        # so it has to be reported instead of wrapping into a bogus dimension.
+        # These stay lazy, so nothing near this size is allocated.
+        big = mx.zeros(2**30)
+        for parts in (3, 4, 5):
+            with self.assertRaises(OverflowError) as cm:
+                mx.concatenate([big] * parts)
+            self.assertIn(str(2**30 * parts), str(cm.exception))
+
+        # repeat and kron multiply a dimension, and used to wrap into a
+        # negative or zero one that only surfaced later as a confusing reshape
+        # error naming a shape the caller never asked for.
+        for parts in (2, 3, 4):
+            with self.assertRaises(OverflowError) as cm:
+                mx.repeat(big, parts)
+            self.assertIn(str(2**30 * parts), str(cm.exception))
+
+        with self.assertRaises(OverflowError) as cm:
+            mx.kron(mx.zeros(2**16), mx.zeros(2**16))
+        self.assertIn(str(2**32), str(cm.exception))
+
         # Negative overflow (< int32 min) is caught too.
         too_negative = -(2**31) - 1
         with self.assertRaises(OverflowError) as cm: