| 345 | self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, c) |
| 346 | |
| 347 | def test_fftshift_errors(self): |
| 348 | # Test invalid axes |
| 349 | x = mx.array(np.random.rand(4, 4).astype(np.float32)) |
| 350 | with self.assertRaises(ValueError): |
| 351 | mx.fft.fftshift(x, axes=[2]) |
| 352 | with self.assertRaises(ValueError): |
| 353 | mx.fft.fftshift(x, axes=[-3]) |
| 354 | |
| 355 | # Test empty array |
| 356 | x = mx.array([]) |
| 357 | self.assertTrue(mx.array_equal(mx.fft.fftshift(x), x)) |
| 358 | |
| 359 | @unittest.skipIf(not has_torch, "requires PyTorch") |
| 360 | def test_fft_grads(self): |