Fix negative dim indexing (#2994)

Co-authored-by: KD2YCU <[email protected]>
Co-authored-by: Awni Hannun <[email protected]>

mlx/array.h

@@ -121,7 +121,7 @@ class array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto shape(int dim) const {
-    return shape().at(dim < 0 ? dim + ndim() : dim);
+    return shape().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
   }
 
   /** The strides of the array. */
@@ -135,7 +135,7 @@ class array {
    *  This function supports negative indexing and provides
    *  bounds checking. */
   auto strides(int dim) const {
-    return strides().at(dim < 0 ? dim + ndim() : dim);
+    return strides().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
   }
 
   /** Get the arrays data type. */

tests/array_tests.cpp

@@ -1,5 +1,8 @@
 // Copyright © 2023 Apple Inc.
+#include <cassert>
 #include <climits>
+#include <stdexcept>
+#include <vector>
 
 #include "doctest/doctest.h"
 
@@ -633,3 +636,28 @@ TEST_CASE("test make array from user buffer") {
   // deleter should always get called
   CHECK_EQ(count, 1);
 }
+
+TEST_CASE("test negative indexing for shape/strides") {
+  // 2D array: shape = {2, 3}
+  std::vector<float> data(6, 1.0f);
+  array a(data.begin(), Shape{2, 3});
+
+  // Valid negative indexing
+  CHECK_EQ(a.shape(-1), a.shape(1));
+  CHECK_EQ(a.shape(-2), a.shape(0));
+  CHECK_EQ(a.shape(-1), 3);
+  CHECK_EQ(a.shape(-2), 2);
+
+  CHECK_EQ(a.strides(-1), a.strides(1));
+  CHECK_EQ(a.strides(-2), a.strides(0));
+  CHECK_EQ(a.strides(-1), 1);
+  CHECK_EQ(a.strides(-2), 3);
+
+  // Invalid: too negative
+  CHECK_THROWS_AS(a.shape(-3), std::out_of_range);
+  CHECK_THROWS_AS(a.strides(-3), std::out_of_range);
+
+  // Invalid: too positive
+  CHECK_THROWS_AS(a.shape(2), std::out_of_range);
+  CHECK_THROWS_AS(a.strides(2), std::out_of_range);
+}