From 34aaf8b708c6a0b2f341a283d2c2a273c3f7ecb1 Mon Sep 17 00:00:00 2001
From: blanky <aaishwarymishra@gmail.com>
Date: Mon, 24 Aug 2026 12:05:41 +0530
Subject: [PATCH 1/2] feat: Added explicit index and bytes

---
 python/src/array.cpp       |  8 ++++++++
 python/src/indexing.cpp    |  5 +++++
 python/tests/test_array.py | 15 +++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/python/src/array.cpp b/python/src/array.cpp
index 3d9e3f7d1d..f701ea5818 100644
--- a/python/src/array.cpp
+++ b/python/src/array.cpp
@@ -1034,6 +1034,14 @@ void init_array(nb::module_& m) {
           [](mx::array& a) {
             return nb::cast<std::complex<double>>(to_scalar(a));
           })
+      .def("__index__", [](mx::array& a) { return nb::int_(to_scalar(a)); })
+      .def(
+          "__bytes__",
+          [](mx::array& a) {
+            a.eval();
+            return nb::bytes(
+                reinterpret_cast<const char*>(a.data<void>()), a.nbytes());
+          })
       .def(
           "__format__",
           [](mx::array& a, nb::object format_spec) {
diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp
index 0e230ba529..0079374849 100644
--- a/python/src/indexing.cpp
+++ b/python/src/indexing.cpp
@@ -5,6 +5,7 @@
 
 #include <nanobind/ndarray.h>
 
+#include "mlx/dtype.h"
 #include "mlx/ops.h"
 #include "python/src/convert.h"
 #include "python/src/indexing.h"
@@ -20,6 +21,10 @@ bool is_index_scalar(const nb::object& obj) {
   if (nb::isinstance<nb::bool_>(obj)) {
     return false;
   }
+  if (nb::isinstance<mx::array>(obj) &&
+      nb::cast<mx::array>(obj).dtype() == mx::bool_) {
+    return false;
+  }
   if (!PyIndex_Check(obj.ptr())) {
     return false;
   }
diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index f82ece534d..58bc079d01 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -140,6 +140,21 @@ def test_scalar_conversion(self):
                 self.assertListEqual(list(z.shape), list(x.shape))
                 self.assertListEqual(list(z.shape), list(y.shape))
 
+    def test_index_conversion(self):
+        for dtype in [
+            mx.uint8,
+            mx.uint16,
+            mx.uint32,
+            mx.uint64,
+            mx.int8,
+            mx.int16,
+            mx.int32,
+            mx.int64,
+        ]:
+            with self.subTest(dtype=dtype):
+                self.assertEqual(operator.index(mx.array(2, dtype)), 2)
+                self.assertEqual(list(range(mx.array(3, dtype))), [0, 1, 2])
+
     def test_finfo(self):
         with self.assertRaises(ValueError):
             mx.finfo(mx.int32)

From 7df74b5f8d5c944a4acde12c868090634f4132cc Mon Sep 17 00:00:00 2001
From: blanky <aaishwarymishra@gmail.com>
Date: Fri, 28 Aug 2026 23:52:03 +0530
Subject: [PATCH 2/2] add explicit dim and type check in index

---
 python/src/array.cpp       | 10 +++++++++-
 python/src/indexing.cpp    |  3 +--
 python/tests/test_array.py |  9 +++++++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/python/src/array.cpp b/python/src/array.cpp
index f701ea5818..ebc0a5d8f6 100644
--- a/python/src/array.cpp
+++ b/python/src/array.cpp
@@ -1034,7 +1034,15 @@ void init_array(nb::module_& m) {
           [](mx::array& a) {
             return nb::cast<std::complex<double>>(to_scalar(a));
           })
-      .def("__index__", [](mx::array& a) { return nb::int_(to_scalar(a)); })
+      .def(
+          "__index__",
+          [](mx::array& a) {
+            if (!mx::issubdtype(a.dtype(), mx::integer) || a.ndim() != 0) {
+              throw nb::type_error(
+                  "Only 0-dimensional integer arrays can be converted to an index.");
+            }
+            return nb::int_(to_scalar(a));
+          })
       .def(
           "__bytes__",
           [](mx::array& a) {
diff --git a/python/src/indexing.cpp b/python/src/indexing.cpp
index 0079374849..5bf4677505 100644
--- a/python/src/indexing.cpp
+++ b/python/src/indexing.cpp
@@ -21,8 +21,7 @@ bool is_index_scalar(const nb::object& obj) {
   if (nb::isinstance<nb::bool_>(obj)) {
     return false;
   }
-  if (nb::isinstance<mx::array>(obj) &&
-      nb::cast<mx::array>(obj).dtype() == mx::bool_) {
+  if (nb::isinstance<mx::array>(obj)) {
     return false;
   }
   if (!PyIndex_Check(obj.ptr())) {
diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index 58bc079d01..17e49eae1d 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -155,6 +155,15 @@ def test_index_conversion(self):
                 self.assertEqual(operator.index(mx.array(2, dtype)), 2)
                 self.assertEqual(list(range(mx.array(3, dtype))), [0, 1, 2])
 
+    def test_index_conversion_invalid(self):
+        for dtype in [mx.float16, mx.float32, mx.bfloat16, mx.complex64, mx.bool_]:
+            with self.subTest(dtype=dtype):
+                with self.assertRaises(TypeError):
+                    operator.index(mx.array(2, dtype))
+
+                with self.assertRaises(TypeError):
+                    list(range(mx.array(3, dtype)))
+
     def test_finfo(self):
         with self.assertRaises(ValueError):
             mx.finfo(mx.int32)