Skip to content

Commit

Permalink
expose itemsize and nbytes as for numpy arrays
Browse files Browse the repository at this point in the history
  • Loading branch information
dastrobu committed Dec 25, 2023
1 parent a123c3c commit 6824099
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 0 deletions.
8 changes: 8 additions & 0 deletions python/src/array.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,14 @@ 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
Expand Down
2 changes: 2 additions & 0 deletions python/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit 6824099

Please sign in to comment.