Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions python/src/array.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1046,9 +1046,10 @@ void init_array(nb::module_& m) {
.def(
"__bytes__",
[](mx::array& a) {
a.eval();
auto c = mx::contiguous(a);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Won't this be better? We might want to save a second call if a is already packed/contiguous :

a.flags().row_contiguous ? a : mx::contiguous(a);

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When a is not evaluated a.flags().row_contiguous would return garbage so we have to use mx::contiguous.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh alright. Then we may skip this. Thanks.

c.eval();
return nb::bytes(
reinterpret_cast<const char*>(a.data<void>()), a.nbytes());
reinterpret_cast<const char*>(c.data<void>()), c.nbytes());
})
.def(
"__format__",
Expand Down
9 changes: 9 additions & 0 deletions python/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -2042,6 +2042,15 @@ def test_buffer_protocol(self):
self.assertEqual(b"aaaaaaaaaa", ab[::2])
self.assertEqual(b"abcdefghij", ab[1::2])

# Test bytes on non-contiguous arrays
a = mx.arange(10, dtype=mx.uint8)
self.assertEqual(bytes(a[::2]), b"\x00\x02\x04\x06\x08")
self.assertEqual(bytes(a[::-1]), b"\x09\x08\x07\x06\x05\x04\x03\x02\x01\x00")
b = mx.arange(6, dtype=mx.int32).reshape(2, 3).T
self.assertEqual(bytes(b), np.array(b).tobytes())
c = mx.broadcast_to(mx.array([1, 2], dtype=mx.uint8), (3, 2))
self.assertEqual(bytes(c), np.array(c).tobytes())

def test_buffer_protocol_ref_counting(self):
a = mx.arange(3)
wr = weakref.ref(a)
Expand Down