From 83bb7891dbd3a3414cd060180f0fafa76905d9f4 Mon Sep 17 00:00:00 2001
From: Dan Anderson <dan@kd2ycu.com>
Date: Tue, 20 Jan 2026 09:24:33 -0500
Subject: [PATCH] Fix negative dim indexing (#2994)

Co-authored-by: KD2YCU <me@kd2ycu.com>
Co-authored-by: Awni Hannun <awni@apple.com>
---
 mlx/array.h           |  4 ++--
 tests/array_tests.cpp | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/mlx/array.h b/mlx/array.h
index 645fa68b..1b7f7ed0 100644
--- a/mlx/array.h
+++ b/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. */
diff --git a/tests/array_tests.cpp b/tests/array_tests.cpp
index 55af2df1..84909472 100644
--- a/tests/array_tests.cpp
+++ b/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);
+}