| 3327 | mx.broadcast_shapes() |
| 3328 | |
| 3329 | def test_sort_nan(self): |
| 3330 | for dtype in [mx.float32, mx.float16, mx.bfloat16]: |
| 3331 | with self.subTest(dtype=dtype): |
| 3332 | x = mx.array([3.0, mx.nan, 2.0, 0.0], dtype=dtype) |
| 3333 | expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype) |
| 3334 | self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) |
| 3335 | |
| 3336 | x = mx.array([3.0 + 1j, mx.nan + 2j, 2.0 + 1j, 0.0 + 1j]) |
| 3337 | expected = mx.array([0.0 + 1j, 2.0 + 1j, 3.0 + 1j, mx.nan + 2j]) |
| 3338 | self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True)) |
| 3339 | |
| 3340 | def test_argsort_nan(self): |
| 3341 | for dtype in [mx.float32, mx.float16, mx.bfloat16]: |