| 1935 | self.assertEqualArray(c_mlx, mx.array(c_np)) |
| 1936 | |
| 1937 | def test_pad(self): |
| 1938 | pad_width_and_values = [ |
| 1939 | ([(1, 1), (1, 1), (1, 1)], 0), |
| 1940 | ([(1, 1), (1, 1), (1, 1)], 5), |
| 1941 | ([(3, 0), (0, 2), (5, 7)], 0), |
| 1942 | ([(3, 0), (0, 2), (5, 7)], -7), |
| 1943 | ([(0, 0), (0, 0), (0, 0)], 0), |
| 1944 | ] |
| 1945 | |
| 1946 | for pw, v in pad_width_and_values: |
| 1947 | with self.subTest(pad_width=pw, value=v): |
| 1948 | a_npy = np.random.randn(16, 16, 16).astype(np.float32) |
| 1949 | a_mlx = mx.array(a_npy) |
| 1950 | |
| 1951 | b_npy = np.pad(a_npy, pw, constant_values=v) |
| 1952 | b_mlx = mx.pad(a_mlx, pw, constant_values=v) |
| 1953 | |
| 1954 | self.assertEqual(list(b_npy.shape), list(b_mlx.shape)) |
| 1955 | self.assertTrue(np.allclose(b_npy, b_mlx, atol=1e-6)) |
| 1956 | |
| 1957 | b_npy = np.pad(a_npy, pw, mode="edge") |
| 1958 | b_mlx = mx.pad(a_mlx, pw, mode="edge") |
| 1959 | |
| 1960 | self.assertEqual(list(b_npy.shape), list(b_mlx.shape)) |
| 1961 | self.assertTrue(np.allclose(b_npy, b_mlx, atol=1e-6)) |
| 1962 | |
| 1963 | a = mx.zeros((1, 1, 1)) |
| 1964 | self.assertEqual(mx.pad(a, 1).shape, (3, 3, 3)) |
| 1965 | self.assertEqual(mx.pad(a, (1,)).shape, (3, 3, 3)) |
| 1966 | self.assertEqual(mx.pad(a, [1]).shape, (3, 3, 3)) |
| 1967 | self.assertEqual(mx.pad(a, (1, 2)).shape, (4, 4, 4)) |
| 1968 | self.assertEqual(mx.pad(a, [(1, 2)]).shape, (4, 4, 4)) |
| 1969 | self.assertEqual(mx.pad(a, ((1, 2),)).shape, (4, 4, 4)) |
| 1970 | self.assertEqual(mx.pad(a, ((1, 2), (2, 1), (2, 2))).shape, (4, 4, 5)) |
| 1971 | |
| 1972 | # Test grads |
| 1973 | a_fwd = mx.array(np.random.rand(16, 16).astype(np.float32)) |
| 1974 | a_bwd = mx.ones((22, 22)) |
| 1975 | f = lambda x: mx.pad(x, ((4, 2), (2, 4))) |
| 1976 | |
| 1977 | _, df = mx.vjp(f, [a_fwd], [a_bwd]) |
| 1978 | self.assertTrue(mx.allclose(a_bwd[4:-2, 2:-4], df[0]).item()) |
| 1979 | |
| 1980 | def test_where(self): |
| 1981 | self.assertCmpNumpy([True, mx.array([[1, 2], [3, 4]]), 1], mx.where, np.where) |