(self)
| 1707 | self.assertTrue(a_np.flags.writeable) |
| 1708 | |
| 1709 | def test_buffer_protocol(self): |
| 1710 | dtypes_list = [ |
| 1711 | (mx.bool_, np.bool_, None), |
| 1712 | (mx.uint8, np.uint8, np.iinfo), |
| 1713 | (mx.uint16, np.uint16, np.iinfo), |
| 1714 | (mx.uint32, np.uint32, np.iinfo), |
| 1715 | (mx.uint64, np.uint64, np.iinfo), |
| 1716 | (mx.int8, np.int8, np.iinfo), |
| 1717 | (mx.int16, np.int16, np.iinfo), |
| 1718 | (mx.int32, np.int32, np.iinfo), |
| 1719 | (mx.int64, np.int64, np.iinfo), |
| 1720 | (mx.float16, np.float16, np.finfo), |
| 1721 | (mx.float32, np.float32, np.finfo), |
| 1722 | (mx.complex64, np.complex64, np.finfo), |
| 1723 | ] |
| 1724 | |
| 1725 | for mlx_dtype, np_dtype, info_fn in dtypes_list: |
| 1726 | a_np = np.random.uniform(low=0, high=100, size=(3, 4)).astype(np_dtype) |
| 1727 | if info_fn is not None: |
| 1728 | info = info_fn(np_dtype) |
| 1729 | a_np[0, 0] = info.min |
| 1730 | a_np[0, 1] = info.max |
| 1731 | a_mx = mx.array(a_np) |
| 1732 | for f in [lambda x: x, lambda x: x.T]: |
| 1733 | mv_mx = memoryview(f(a_mx)) |
| 1734 | mv_np = memoryview(f(a_np)) |
| 1735 | self.assertEqual(mv_mx.strides, mv_np.strides, f"{mlx_dtype}{np_dtype}") |
| 1736 | self.assertEqual(mv_mx.shape, mv_np.shape, f"{mlx_dtype}{np_dtype}") |
| 1737 | # correct buffer format for 8 byte (unsigned) 'long long' is Q/q, see |
| 1738 | # https://docs.python.org/3.10/library/struct.html#format-characters |
| 1739 | # numpy returns L/l, as 'long' is equivalent to 'long long' on 64bit machines, so q and l are equivalent |
| 1740 | # see https://github.com/pybind/pybind11/issues/1908 |
| 1741 | if np_dtype == np.uint64: |
| 1742 | self.assertEqual(mv_mx.format, "Q", f"{mlx_dtype}{np_dtype}") |
| 1743 | elif np_dtype == np.int64: |
| 1744 | self.assertEqual(mv_mx.format, "q", f"{mlx_dtype}{np_dtype}") |
| 1745 | # for windows long is 32bit and numpy returns L/l. |
| 1746 | elif np_dtype == np.uint32 and platform.system() == "Windows": |
| 1747 | self.assertEqual(mv_mx.format, "I", f"{mlx_dtype}{np_dtype}") |
| 1748 | elif np_dtype == np.int32 and platform.system() == "Windows": |
| 1749 | self.assertEqual(mv_mx.format, "i", f"{mlx_dtype}{np_dtype}") |
| 1750 | else: |
| 1751 | self.assertEqual( |
| 1752 | mv_mx.format, mv_np.format, f"{mlx_dtype}{np_dtype}" |
| 1753 | ) |
| 1754 | self.assertFalse(mv_mx.readonly) |
| 1755 | back_to_npy = np.array(mv_mx, copy=False) |
| 1756 | self.assertEqualArray( |
| 1757 | back_to_npy, |
| 1758 | f(a_np), |
| 1759 | atol=0, |
| 1760 | rtol=0, |
| 1761 | ) |
| 1762 | |
| 1763 | # extra test for bfloat16, which is not numpy convertible |
| 1764 | a_mx = mx.random.uniform(low=0, high=100, shape=(3, 4), dtype=mx.bfloat16) |
| 1765 | mv_mx = memoryview(a_mx) |
| 1766 | self.assertEqual(mv_mx.strides, (8, 2)) |
nothing calls this directly
no test coverage detected