Fix buffer protocol buffer size designation (#1010)

This commit is contained in:
Angelos Katharopoulos
2024-04-19 06:06:13 -07:00
committed by GitHub
parent 090ff659dc
commit ef5f7d1aea
2 changed files with 13 additions and 1 deletions

View File

@@ -2,6 +2,7 @@
import operator
import pickle
import sys
import unittest
import weakref
from copy import copy, deepcopy
@@ -1497,6 +1498,17 @@ class TestArray(mlx_tests.MLXTestCase):
e = cm.exception
self.assertTrue("Item size 2 for PEP 3118 buffer format string" in str(e))
# Test buffer protocol with non-arrays ie bytes
a = ord("a") * 257 + mx.arange(10).astype(mx.int16)
ab = bytes(a)
self.assertEqual(len(ab), 20)
if sys.byteorder == "little":
self.assertEqual(b"aaaaaaaaaa", ab[1::2])
self.assertEqual(b"abcdefghij", ab[::2])
else:
self.assertEqual(b"aaaaaaaaaa", ab[::2])
self.assertEqual(b"abcdefghij", ab[1::2])
def test_buffer_protocol_ref_counting(self):
a = mx.arange(3)
wr = weakref.ref(a)