| 2191 | self.assertEqual(mx.expand_dims(a, [0, -1]).shape, (1, 2, 2, 1)) |
| 2192 | |
| 2193 | def test_sort(self): |
| 2194 | shape = (6, 4, 10) |
| 2195 | tests = product( |
| 2196 | ("int32", "float32", "complex64"), # type |
| 2197 | (None, 0, 1, 2), # axis |
| 2198 | (True, False), # strided |
| 2199 | ) |
| 2200 | for dtype, axis, strided in tests: |
| 2201 | with self.subTest(dtype=dtype, axis=axis, strided=strided): |
| 2202 | np.random.seed(0) |
| 2203 | np_dtype = getattr(np, dtype) |
| 2204 | if np.issubdtype(np_dtype, np.complexfloating): |
| 2205 | a_np = ( |
| 2206 | np.random.uniform(0, 100, size=shape) |
| 2207 | + 1j * np.random.uniform(0, 100, size=shape) |
| 2208 | ).astype(np_dtype) |
| 2209 | else: |
| 2210 | a_np = np.random.uniform(0, 100, size=shape).astype(np_dtype) |
| 2211 | a_mx = mx.array(a_np) |
| 2212 | if strided: |
| 2213 | a_mx = a_mx[::2, :, ::2] |
| 2214 | a_np = a_np[::2, :, ::2] |
| 2215 | |
| 2216 | b_np = np.sort(a_np, axis=axis) |
| 2217 | b_mx = mx.sort(a_mx, axis=axis) |
| 2218 | |
| 2219 | self.assertTrue(np.array_equal(b_np, b_mx)) |
| 2220 | self.assertEqual(b_mx.dtype, a_mx.dtype) |
| 2221 | |
| 2222 | c_np = np.argsort(a_np, axis=axis) |
| 2223 | c_mx = mx.argsort(a_mx, axis=axis) |
| 2224 | d_np = np.take_along_axis(a_np, c_np, axis=axis) |
| 2225 | d_mx = mx.take_along_axis(a_mx, c_mx, axis=axis) |
| 2226 | |
| 2227 | self.assertTrue(np.array_equal(d_np, d_mx)) |
| 2228 | self.assertEqual(c_mx.dtype, mx.uint32) |
| 2229 | |
| 2230 | # Set random seed |
| 2231 | np.random.seed(0) |
| 2232 | |
| 2233 | # Test multi-block sort |
| 2234 | for strided in (False, True): |
| 2235 | with self.subTest(strided=strided): |
| 2236 | a_np = np.random.normal(size=(32769,)).astype(np.float32) |
| 2237 | a_mx = mx.array(a_np) |
| 2238 | |
| 2239 | if strided: |
| 2240 | a_mx = a_mx[::3] |
| 2241 | a_np = a_np[::3] |
| 2242 | |
| 2243 | b_np = np.sort(a_np) |
| 2244 | b_mx = mx.sort(a_mx) |
| 2245 | |
| 2246 | self.assertTrue(np.array_equal(b_np, b_mx)) |
| 2247 | self.assertEqual(b_mx.dtype, a_mx.dtype) |
| 2248 | |
| 2249 | # Test multi-dum multi-block sort |
| 2250 | a_np = np.random.normal(size=(2, 4, 32769)).astype(np.float32) |