| 1585 | self.assertTrue(np.allclose(expected, np.array(out), equal_nan=True)) |
| 1586 | |
| 1587 | def test_trig_ops(self): |
| 1588 | def test_ops(npop, mlxop, x, y, atol, rtol): |
| 1589 | r_np = npop(x) |
| 1590 | r_mlx = mlxop(y) |
| 1591 | mx.eval(r_mlx) |
| 1592 | |
| 1593 | self.assertTrue( |
| 1594 | np.allclose(r_np, r_mlx, atol=atol, rtol=rtol, equal_nan=True) |
| 1595 | ) |
| 1596 | |
| 1597 | x = np.random.rand(9, 12, 18) |
| 1598 | xi = np.random.rand(9, 12, 18) |
| 1599 | base_ops = ["sin", "cos", "tan"] |
| 1600 | hyperbolic_ops = ["sinh", "cosh", "tanh"] |
| 1601 | all_fwd_ops = base_ops + hyperbolic_ops |
| 1602 | |
| 1603 | for op in all_fwd_ops: |
| 1604 | with self.subTest(op=op): |
| 1605 | float_dtypes = [("float16", 1e-3, 1e-3), ("float32", 1e-6, 1e-5)] |
| 1606 | |
| 1607 | for dtype, atol, rtol in float_dtypes: |
| 1608 | with self.subTest(dtype=dtype): |
| 1609 | x_ = x.astype(getattr(np, dtype)) |
| 1610 | y_ = mx.array(x_) |
| 1611 | test_ops(getattr(np, op), getattr(mx, op), x_, y_, atol, rtol) |
| 1612 | |
| 1613 | with self.subTest(op=op): |
| 1614 | dtype = "complex64" |
| 1615 | with self.subTest(dtype=dtype): |
| 1616 | x_ = x + 1.0j * xi |
| 1617 | x_ = x_.astype(getattr(np, dtype)) |
| 1618 | y_ = mx.array(x_) |
| 1619 | test_ops(getattr(np, op), getattr(mx, op), x_, y_, 1e-5, 1e-5) |
| 1620 | |
| 1621 | with self.subTest(op="arc" + op): |
| 1622 | float_dtypes = [("float16", 1e-3, 1e-3), ("float32", 1e-6, 1e-5)] |
| 1623 | op_inv = "arc" + op |
| 1624 | |
| 1625 | for dtype, atol, rtol in float_dtypes: |
| 1626 | with self.subTest(dtype=dtype): |
| 1627 | np_op_fwd = getattr(np, op) |
| 1628 | x_ = np_op_fwd(x).astype(getattr(np, dtype)) |
| 1629 | y_ = mx.array(x_) |
| 1630 | test_ops( |
| 1631 | getattr(np, op_inv), getattr(mx, op_inv), x_, y_, atol, rtol |
| 1632 | ) |
| 1633 | |
| 1634 | # Test grads |
| 1635 | np_vjp_funcs = { |
| 1636 | "sin": lambda primal, cotan: cotan * np.cos(primal), |
| 1637 | "cos": lambda primal, cotan: -cotan * np.sin(primal), |
| 1638 | "tan": lambda primal, cotan: cotan / (np.cos(primal) ** 2), |
| 1639 | "sinh": lambda primal, cotan: cotan * np.cosh(primal), |
| 1640 | "cosh": lambda primal, cotan: cotan * np.sinh(primal), |
| 1641 | "tanh": lambda primal, cotan: cotan / (np.cosh(primal) ** 2), |
| 1642 | "arcsin": lambda primal, cotan: cotan / np.sqrt(1.0 - primal**2), |
| 1643 | "arccos": lambda primal, cotan: -cotan / np.sqrt(1.0 - primal**2), |
| 1644 | "arctan": lambda primal, cotan: cotan / (1.0 + primal**2), |