(self)
| 2066 | self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3)) |
| 2067 | |
| 2068 | def test_scans(self): |
| 2069 | a_npy = np.random.randn(32, 32, 32).astype(np.float32) |
| 2070 | a_mlx = mx.array(a_npy) |
| 2071 | |
| 2072 | for op in ["cumsum", "cumprod"]: |
| 2073 | npop = getattr(np, op) |
| 2074 | mxop = getattr(mx, op) |
| 2075 | for axis in (None, 0, 1, 2): |
| 2076 | c_npy = npop(a_npy, axis=axis) |
| 2077 | c_mlx = mxop(a_mlx, axis=axis) |
| 2078 | self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3)) |
| 2079 | |
| 2080 | # Complex test |
| 2081 | |
| 2082 | a_npy = np.random.randn(32, 32, 32).astype(np.float32) + 0.5j |
| 2083 | a_mlx = mx.array(a_npy) |
| 2084 | |
| 2085 | for op in ["cumsum", "cumprod"]: |
| 2086 | npop = getattr(np, op) |
| 2087 | mxop = getattr(mx, op) |
| 2088 | for axis in (None, 0, 1, 2): |
| 2089 | c_npy = npop(a_npy, axis=axis) |
| 2090 | c_mlx = mxop(a_mlx, axis=axis) |
| 2091 | self.assertTrue(np.allclose(c_npy, c_mlx, rtol=1e-3, atol=1e-3)) |
| 2092 | |
| 2093 | a_mlx = mx.random.randint(shape=(32, 32, 32), low=-100, high=100) |
| 2094 | for dt in [mx.int32, mx.int64]: |
| 2095 | mxx = a_mlx.astype(dt) |
| 2096 | npx = np.array(mxx) |
| 2097 | for op in ["cumsum", "cumprod"]: |
| 2098 | npop = getattr(np, op) |
| 2099 | mxop = getattr(mx, op) |
| 2100 | for axis in (None, 0, 1, 2): |
| 2101 | c_npy = npop(npx, axis=axis, dtype=npx.dtype) |
| 2102 | c_mlx = mxop(mxx, axis=axis) |
| 2103 | self.assertTrue(np.array_equal(c_npy, c_mlx)) |
| 2104 | |
| 2105 | a_mlx = mx.random.randint(shape=(32, 32, 32), low=-100, high=100) |
| 2106 | for op in ["cumsum", "cumprod", "cummax", "cummin"]: |
| 2107 | mxop = getattr(mx, op) |
| 2108 | c1 = mxop(a_mlx, axis=2) |
| 2109 | c2 = mxop(a_mlx, axis=2, inclusive=False, reverse=False) |
| 2110 | self.assertTrue(mx.array_equal(c1[:, :, :-1], c2[:, :, 1:])) |
| 2111 | c1 = mxop(a_mlx, axis=1) |
| 2112 | c2 = mxop(a_mlx, axis=1, inclusive=False, reverse=False) |
| 2113 | self.assertTrue(mx.array_equal(c1[:, :-1, :], c2[:, 1:, :])) |
| 2114 | c1 = mxop(a_mlx, axis=0) |
| 2115 | c2 = mxop(a_mlx, axis=0, inclusive=False, reverse=False) |
| 2116 | self.assertTrue(mx.array_equal(c1[:-1, :, :], c2[1:, :, :])) |
| 2117 | |
| 2118 | rev_idx = mx.arange(31, -1, -1) |
| 2119 | c1 = mxop(a_mlx[:, :, rev_idx], axis=2)[:, :, rev_idx] |
| 2120 | c2 = mxop(a_mlx, axis=2, inclusive=True, reverse=True) |
| 2121 | self.assertTrue(mx.array_equal(c1, c2)) |
| 2122 | c1 = mxop(a_mlx[:, rev_idx, :], axis=1)[:, rev_idx, :] |
| 2123 | c2 = mxop(a_mlx, axis=1, inclusive=True, reverse=True) |
| 2124 | self.assertTrue(mx.array_equal(c1, c2)) |
| 2125 | c1 = mxop(a_mlx[rev_idx, :, :], axis=0)[rev_idx, :, :] |
nothing calls this directly
no test coverage detected