| 584 | self.assertListEqual(mx.transpose(x).tolist(), expected) |
| 585 | |
| 586 | def test_transpose_axis(self): |
| 587 | x = mx.array( |
| 588 | [ |
| 589 | [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], |
| 590 | [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]], |
| 591 | ] |
| 592 | ) |
| 593 | expected = [ |
| 594 | [[0, 4, 8], [1, 5, 9], [2, 6, 10], [3, 7, 11]], |
| 595 | [[12, 16, 20], [13, 17, 21], [14, 18, 22], [15, 19, 23]], |
| 596 | ] |
| 597 | |
| 598 | self.assertListEqual(mx.transpose(x, axes=(0, 2, 1)).tolist(), expected) |
| 599 | |
| 600 | def test_move_swap_axes(self): |
| 601 | x = mx.zeros((2, 3, 4)) |