| 1749 | ) |
| 1750 | |
| 1751 | def test_irregular_binary_ops(self): |
| 1752 | # Check transposed binary ops |
| 1753 | dims = [2, 3, 4, 5] |
| 1754 | size = 3 |
| 1755 | trial_mul = 2 |
| 1756 | np.random.seed(0) |
| 1757 | for d in dims: |
| 1758 | anp = np.random.randint(-20, 20, (size**d,)).reshape([size] * d) |
| 1759 | bnp = np.random.randint(-20, 20, (size**d,)).reshape([size] * d) |
| 1760 | for _ in range(trial_mul * d): |
| 1761 | amlx = mx.array(anp) |
| 1762 | bmlx = mx.array(bnp) |
| 1763 | a_t = np.random.permutation(d).tolist() |
| 1764 | b_t = np.random.permutation(d).tolist() |
| 1765 | outnp = np.add(anp.transpose(a_t), bnp.transpose(b_t)) |
| 1766 | outmlx = mx.add(mx.transpose(amlx, a_t), mx.transpose(bmlx, b_t)) |
| 1767 | self.assertTrue(np.array_equal(outnp, outmlx)) |
| 1768 | |
| 1769 | # Check broadcast binary ops |
| 1770 | for d in dims: |
| 1771 | anp = np.random.randint(-20, 20, (size**d,)).reshape([size] * d) |
| 1772 | for n_bsx in range(d): |
| 1773 | bnp = np.random.randint(-20, 20, (size**n_bsx,)).reshape([size] * n_bsx) |
| 1774 | for _ in range(trial_mul * d): |
| 1775 | amlx = mx.array(anp) |
| 1776 | bmlx = mx.array(bnp) |
| 1777 | b_shape = [1] * (d - n_bsx) + [size] * n_bsx |
| 1778 | np.random.shuffle(b_shape) |
| 1779 | outnp = np.add(anp, bnp.reshape(b_shape)) |
| 1780 | outmlx = mx.add(amlx, mx.reshape(bmlx, b_shape)) |
| 1781 | self.assertTrue(np.array_equal(outnp, outmlx)) |
| 1782 | |
| 1783 | # Check strided binary ops |
| 1784 | for d in dims: |
| 1785 | a = np.random.randint(-20, 20, (10,) * d) |
| 1786 | b = np.random.randint(-20, 20, (10,) * d) |
| 1787 | a_ = mx.array(a) |
| 1788 | b_ = mx.array(b) |
| 1789 | for t in permutations(range(d)): |
| 1790 | for s in range(d): |
| 1791 | idx = tuple( |
| 1792 | [slice(None)] * s |
| 1793 | + [slice(None, None, 2)] |
| 1794 | + [slice(None)] * (d - s - 1) |
| 1795 | ) |
| 1796 | c = a.transpose(t)[idx] + b[idx] |
| 1797 | c_ = mx.transpose(a_, t)[idx] + b_[idx] |
| 1798 | self.assertTrue(np.array_equal(c, c_)) |
| 1799 | |
| 1800 | def test_softmax(self): |
| 1801 | cases = [(np.float32, 1e-6), (np.float16, 1e-3)] |