Add printoptions (#3333)
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user