| 1618 | self.assertTrue(np.array_equal(b_np, b_mx)) |
| 1619 | |
| 1620 | def test_api(self): |
| 1621 | x = mx.array(np.random.rand(10, 10, 10)) |
| 1622 | ops = [ |
| 1623 | ("reshape", (100, -1)), |
| 1624 | "square", |
| 1625 | "sqrt", |
| 1626 | "rsqrt", |
| 1627 | "reciprocal", |
| 1628 | "exp", |
| 1629 | "log", |
| 1630 | "sin", |
| 1631 | "cos", |
| 1632 | "log1p", |
| 1633 | "abs", |
| 1634 | "log10", |
| 1635 | "log2", |
| 1636 | "conj", |
| 1637 | ("all", 1), |
| 1638 | ("any", 1), |
| 1639 | ("transpose", (0, 2, 1)), |
| 1640 | ("sum", 1), |
| 1641 | ("prod", 1), |
| 1642 | ("min", 1), |
| 1643 | ("max", 1), |
| 1644 | ("logcumsumexp", 1), |
| 1645 | ("logsumexp", 1), |
| 1646 | ("mean", 1), |
| 1647 | ("var", 1), |
| 1648 | ("argmin", 1), |
| 1649 | ("argmax", 1), |
| 1650 | ("cummax", 1), |
| 1651 | ("cummin", 1), |
| 1652 | ("cumprod", 1), |
| 1653 | ("cumsum", 1), |
| 1654 | ("diagonal", 0, 0, 1), |
| 1655 | ("flatten", 0, -1), |
| 1656 | ("moveaxis", 1, 2), |
| 1657 | ("round", 2), |
| 1658 | ("std", 1, True, 0), |
| 1659 | ("swapaxes", 1, 2), |
| 1660 | ] |
| 1661 | for op in ops: |
| 1662 | if isinstance(op, tuple): |
| 1663 | op, *args = op |
| 1664 | else: |
| 1665 | args = tuple() |
| 1666 | y1 = getattr(mx, op)(x, *args) |
| 1667 | y2 = getattr(x, op)(*args) |
| 1668 | self.assertEqual(y1.dtype, y2.dtype) |
| 1669 | self.assertEqual(y1.shape, y2.shape) |
| 1670 | self.assertTrue(mx.array_equal(y1, y2)) |
| 1671 | |
| 1672 | y1 = mx.split(x, 2) |
| 1673 | y2 = x.split(2) |
| 1674 | self.assertEqual(len(y1), 2) |
| 1675 | self.assertEqual(len(y1), len(y2)) |
| 1676 | self.assertTrue(mx.array_equal(y1[0], y2[0])) |
| 1677 | self.assertTrue(mx.array_equal(y1[1], y2[1])) |