| 535 | self.assertEqual(result.tolist(), [False, False, False]) |
| 536 | |
| 537 | def test_round(self): |
| 538 | # float |
| 539 | x = mx.array( |
| 540 | [0.5, -0.5, 1.5, -1.5, -21.03, 19.98, -27, 9, 0.0, -np.inf, np.inf] |
| 541 | ) |
| 542 | expected = [0, -0, 2, -2, -21, 20, -27, 9, 0, -np.inf, np.inf] |
| 543 | self.assertListEqual(mx.round(x).tolist(), expected) |
| 544 | |
| 545 | # complex |
| 546 | y = mx.round(mx.array([22.2 + 3.6j, 18.5 + 98.2j])) |
| 547 | self.assertListEqual(y.tolist(), [22 + 4j, 18 + 98j]) |
| 548 | |
| 549 | # decimals |
| 550 | y0 = mx.round(mx.array([15, 122], mx.int32), decimals=0) |
| 551 | y1 = mx.round(mx.array([15, 122], mx.int32), decimals=-1) |
| 552 | y2 = mx.round(mx.array([15, 122], mx.int32), decimals=-2) |
| 553 | self.assertEqual(y0.dtype, mx.int32) |
| 554 | self.assertEqual(y1.dtype, mx.int32) |
| 555 | self.assertEqual(y2.dtype, mx.int32) |
| 556 | self.assertListEqual(y0.tolist(), [15, 122]) |
| 557 | self.assertListEqual(y1.tolist(), [20, 120]) |
| 558 | self.assertListEqual(y2.tolist(), [0, 100]) |
| 559 | |
| 560 | y1 = mx.round(mx.array([1.537, 1.471], mx.float32), decimals=1) |
| 561 | y2 = mx.round(mx.array([1.537, 1.471], mx.float32), decimals=2) |
| 562 | self.assertTrue(mx.allclose(y1, mx.array([1.5, 1.5]))) |
| 563 | self.assertTrue(mx.allclose(y2, mx.array([1.54, 1.47]))) |
| 564 | |
| 565 | # check round to nearest for different types |
| 566 | dtypes = [mx.bfloat16, mx.float16, mx.float32] |
| 567 | for dtype in dtypes: |
| 568 | x = mx.arange(10, dtype=dtype) - 4.5 |
| 569 | x = mx.round(x) |
| 570 | self.assertEqual( |
| 571 | x.astype(mx.float32).tolist(), |
| 572 | [-4.0, -4.0, -2.0, -2.0, -0.0, 0.0, 2.0, 2.0, 4.0, 4.0], |
| 573 | ) |
| 574 | |
| 575 | def test_transpose_noargs(self): |
| 576 | x = mx.array([[0, 1, 1], [1, 0, 0]]) |