| 1680 | test_ops(np_vjp, mx_vjp, x_, y_, 1e-5, 1e-5) |
| 1681 | |
| 1682 | def test_binary_ops(self): |
| 1683 | def test_ops(npop, mlxop, x1, x2, y1, y2, atol): |
| 1684 | r_np = npop(x1, x2) |
| 1685 | r_mlx = mlxop(y1, y2) |
| 1686 | mx.eval(r_mlx) |
| 1687 | self.assertTrue(np.allclose(r_np, r_mlx, atol=atol)) |
| 1688 | |
| 1689 | r_np = npop(x1[:1], x2) |
| 1690 | r_mlx = mlxop(y1[:1], y2) |
| 1691 | mx.eval(r_mlx) |
| 1692 | self.assertTrue(np.allclose(r_np, r_mlx, atol=atol)) |
| 1693 | |
| 1694 | r_np = npop(x1[:, :1], x2) |
| 1695 | r_mlx = mlxop(y1[:, :1], y2) |
| 1696 | mx.eval(r_mlx) |
| 1697 | self.assertTrue(np.allclose(r_np, r_mlx, atol=atol)) |
| 1698 | |
| 1699 | r_np = npop(x1[:, :, :1], x2) |
| 1700 | r_mlx = mlxop(y1[:, :, :1], y2) |
| 1701 | mx.eval(r_mlx) |
| 1702 | self.assertTrue(np.allclose(r_np, r_mlx, atol=atol)) |
| 1703 | |
| 1704 | x1 = np.maximum(np.random.rand(18, 28, 38), 0.1) |
| 1705 | x2 = np.maximum(np.random.rand(18, 28, 38), 0.1) |
| 1706 | y1 = mx.array(x1) |
| 1707 | y2 = mx.array(x2) |
| 1708 | mx.eval(y1, y2) |
| 1709 | for op in [ |
| 1710 | "add", |
| 1711 | "subtract", |
| 1712 | "multiply", |
| 1713 | "divide", |
| 1714 | "floor_divide", |
| 1715 | "maximum", |
| 1716 | "minimum", |
| 1717 | "power", |
| 1718 | ]: |
| 1719 | with self.subTest(op=op): |
| 1720 | int_dtypes = [ |
| 1721 | "int8", |
| 1722 | "int16", |
| 1723 | "int32", |
| 1724 | "int64", |
| 1725 | "uint8", |
| 1726 | "uint16", |
| 1727 | "uint32", |
| 1728 | "uint64", |
| 1729 | ] |
| 1730 | float_dtypes = ["float16", "float32"] |
| 1731 | |
| 1732 | dtypes = { |
| 1733 | "divide": float_dtypes, |
| 1734 | "power": float_dtypes, |
| 1735 | "floor_divide": ["float32"] + int_dtypes, |
| 1736 | } |
| 1737 | dtypes = dtypes.get(op, int_dtypes + float_dtypes) |
| 1738 | |
| 1739 | for dtype in dtypes: |