| 2622 | self.assertTrue(mx.array_equal(out, mx.zeros((10, 10)))) |
| 2623 | |
| 2624 | def test_diagonal(self): |
| 2625 | x = mx.array( |
| 2626 | [ |
| 2627 | [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], |
| 2628 | [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], |
| 2629 | ] |
| 2630 | ) |
| 2631 | expected = [[0, 13], [4, 17], [8, 21]] |
| 2632 | |
| 2633 | self.assertListEqual(mx.diagonal(x, 0, -1, 0).tolist(), expected) |
| 2634 | |
| 2635 | expected = [[1, 14], [5, 18], [9, 22]] |
| 2636 | self.assertListEqual(mx.diagonal(x, -1, 2, 0).tolist(), expected) |
| 2637 | |
| 2638 | def test_diag(self): |
| 2639 | # Test 1D input |