L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[CUDA] Add FFT support (#3243)

Browse Source
This commit is contained in:
Lucas Newman
2026-03-14 21:02:19 +09:00
committed by GitHub
parent b0564a9112
commit 5d1700493a
9 changed files with 498 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/cuda_skip.py
-10
View File
@@ -6,16 +6,6 @@ cuda_skip = {
"TestBlas.test_gather_matmul",
"TestBlas.test_gather_matmul_grad",
"TestBlas.test_gather_mm_sorted_vjp",
# FFTs NYI
"TestFFT.test_fft",
"TestFFT.test_fft_big_powers_of_two",
"TestFFT.test_fft_contiguity",
"TestFFT.test_fft_exhaustive",
"TestFFT.test_fft_grads",
"TestFFT.test_fft_into_ifft",
"TestFFT.test_fft_large_numbers",
"TestFFT.test_fft_shared_mem",
"TestFFT.test_fftn",
# Lapack ops NYI
"TestLinalg.test_cholesky",
"TestLinalg.test_cholesky_inv",
python/tests/test_fft.py
+11
View File
@@ -91,6 +91,17 @@ class TestFFT(mlx_tests.MLXTestCase):
np_op = getattr(np.fft, op)
self.check_mx_np(mx_op, np_op, x, axes=ax, s=s)
# Explicitly exercise transposed layouts and axes that are not
# physically last in memory order.
xt = np.transpose(a, (1, 2, 0))
self.check_mx_np(mx.fft.fftn, np.fft.fftn, xt, axes=(2, 0))
self.check_mx_np(mx.fft.ifftn, np.fft.ifftn, xt, axes=(2, 0))
rt = np.transpose(r, (1, 2, 0))
self.check_mx_np(mx.fft.rfftn, np.fft.rfftn, rt, axes=(2, 0))
irfft_in = np.ascontiguousarray(np.fft.rfftn(rt, axes=(2, 0)))
self.check_mx_np(mx.fft.irfftn, np.fft.irfftn, irfft_in, axes=(2, 0))
def _run_ffts(self, shape, atol=1e-4, rtol=1e-4):
np.random.seed(9)