From 026c8df81ed7dd8585059d849c8369354af45647 Mon Sep 17 00:00:00 2001
From: Daniel Strobusch <1847260+dastrobu@users.noreply.github.com>
Date: Mon, 25 Dec 2023 17:59:36 +0100
Subject: [PATCH] expose itemsize and nbytes as for numpy arrays]

see:
  * https://numpy.org/doc/stable/reference/generated/numpy.ndarray.nbytes.html
  * https://numpy.org/doc/stable/reference/generated/numpy.ndarray.itemsize.html

relates to https://github.com/ml-explore/mlx-examples/pull/174
---
 python/src/array.cpp       | 4 ++++
 python/tests/test_array.py | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/python/src/array.cpp b/python/src/array.cpp
index 74223cdda4..d5dd00bd81 100644
--- a/python/src/array.cpp
+++ b/python/src/array.cpp
@@ -510,6 +510,10 @@ void init_array(py::module_& m) {
           "size", &array::size, R"pbdoc(Number of elments in the array.)pbdoc")
       .def_property_readonly(
           "ndim", &array::ndim, R"pbdoc(The array's dimension.)pbdoc")
+      .def_property_readonly(
+          "itemsize", &array::itemsize, R"pbdoc(The size of the array's datatype in bytes.)pbdoc")
+      .def_property_readonly(
+          "nbytes", &array::nbytes, R"pbdoc(The number of bytes in the array.)pbdoc")
       // TODO, this makes a deep copy of the shape
       // implement alternatives to use reference
       // https://pybind11.readthedocs.io/en/stable/advanced/cast/stl.html
diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index fb6a24cbca..847d6d1423 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -84,6 +84,8 @@ def test_array_basics(self):
         x = mx.array(1)
         self.assertEqual(x.size, 1)
         self.assertEqual(x.ndim, 0)
+        self.assertEqual(x.itemsize, 4)
+        self.assertEqual(x.nbytes, 4)
         self.assertEqual(x.shape, [])
         self.assertEqual(x.dtype, mx.int32)
         self.assertEqual(x.item(), 1)