(self)
| 1872 | self.assertTrue(mx.array_equal(out, b)) |
| 1873 | |
| 1874 | def test_meshgrid(self): |
| 1875 | x = mx.array([1, 2, 3], dtype=mx.int32) |
| 1876 | y = np.array([1, 2, 3], dtype=np.int32) |
| 1877 | |
| 1878 | # Test single input |
| 1879 | a_mlx = mx.meshgrid(x) |
| 1880 | a_np = np.meshgrid(y) |
| 1881 | self.assertEqualArray(a_mlx[0], mx.array(a_np[0])) |
| 1882 | |
| 1883 | # Test sparse |
| 1884 | a_mlx, b_mlx, c_mlx = mx.meshgrid(x, x, x, sparse=True) |
| 1885 | a_np, b_np, c_np = np.meshgrid(y, y, y, sparse=True) |
| 1886 | self.assertEqualArray(a_mlx, mx.array(a_np)) |
| 1887 | self.assertEqualArray(b_mlx, mx.array(b_np)) |
| 1888 | self.assertEqualArray(c_mlx, mx.array(c_np)) |
| 1889 | |
| 1890 | # Test different lengths |
| 1891 | x = mx.array([1, 2], dtype=mx.int32) |
| 1892 | y = mx.array([1, 2, 3], dtype=mx.int32) |
| 1893 | z = np.array([1, 2], dtype=np.int32) |
| 1894 | w = np.array([1, 2, 3], dtype=np.int32) |
| 1895 | a_mlx, b_mlx = mx.meshgrid(x, y) |
| 1896 | a_np, b_np = np.meshgrid(z, w) |
| 1897 | self.assertEqualArray(a_mlx, mx.array(a_np)) |
| 1898 | self.assertEqualArray(b_mlx, mx.array(b_np)) |
| 1899 | |
| 1900 | # Test empty input |
| 1901 | x = mx.array([], dtype=mx.int32) |
| 1902 | y = np.array([], dtype=np.int32) |
| 1903 | a_mlx = mx.meshgrid(x) |
| 1904 | a_np = np.meshgrid(y) |
| 1905 | self.assertEqualArray(a_mlx[0], mx.array(a_np[0])) |
| 1906 | |
| 1907 | # Test float32 input |
| 1908 | x = mx.array([1.1, 2.2, 3.3], dtype=mx.float32) |
| 1909 | y = np.array([1.1, 2.2, 3.3], dtype=np.float32) |
| 1910 | a_mlx = mx.meshgrid(x, x, x) |
| 1911 | a_np = np.meshgrid(y, y, y) |
| 1912 | self.assertEqualArray(a_mlx[0], mx.array(a_np[0])) |
| 1913 | self.assertEqualArray(a_mlx[1], mx.array(a_np[1])) |
| 1914 | self.assertEqualArray(a_mlx[2], mx.array(a_np[2])) |
| 1915 | |
| 1916 | # Test ij indexing |
| 1917 | x = mx.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=mx.float32) |
| 1918 | y = np.array([1.1, 2.2, 3.3, 4.4, 5.5], dtype=np.float32) |
| 1919 | a_mlx = mx.meshgrid(x, x, indexing="ij") |
| 1920 | a_np = np.meshgrid(y, y, indexing="ij") |
| 1921 | self.assertEqualArray(a_mlx[0], mx.array(a_np[0])) |
| 1922 | self.assertEqualArray(a_mlx[1], mx.array(a_np[1])) |
| 1923 | |
| 1924 | # Test different lengths, sparse, and ij indexing |
| 1925 | a = mx.array([1, 2], dtype=mx.int64) |
| 1926 | b = mx.array([1, 2, 3], dtype=mx.int64) |
| 1927 | c = mx.array([1, 2, 3, 4], dtype=mx.int64) |
| 1928 | x = np.array([1, 2], dtype=np.int64) |
| 1929 | y = np.array([1, 2, 3], dtype=np.int64) |
| 1930 | z = np.array([1, 2, 3, 4], dtype=np.int64) |
| 1931 | a_mlx, b_mlx, c_mlx = mx.meshgrid(a, b, c, sparse=True, indexing="ij") |
nothing calls this directly
no test coverage detected