Add printoptions (#3333)

This commit is contained in:
Christophe Prat
2026-04-02 07:24:48 +02:00
committed by GitHub
parent 80a1c206f9
commit befe42d303
9 changed files with 174 additions and 13 deletions
+20
View File
@@ -597,6 +597,26 @@ class TestArray(mlx_tests.MLXTestCase):
x = mx.array([1 - 1j], dtype=mx.complex64)
expected = "array([1-1j], dtype=complex64)"
def test_array_repr_precision(self):
x = mx.array([1.123456789], dtype=mx.float32)
expected = "array([1.12346], dtype=float32)"
self.assertEqual(str(x), expected)
with mx.printoptions(precision=4):
expected = "array([1.1235], dtype=float32)"
self.assertEqual(str(x), expected)
mx.set_printoptions(precision=2)
expected = "array([1.12], dtype=float32)"
self.assertEqual(str(x), expected)
x = mx.sin(x)
expected = "array([0.90], dtype=float32)"
self.assertEqual(str(x), expected)
with mx.printoptions(precision=4):
expected = "array([0.9016], dtype=float32)"
self.assertEqual(str(x), expected)
def test_array_to_list(self):
types = [mx.bool_, mx.uint32, mx.int32, mx.int64, mx.float32]
for t in types: