mlx_rs/fft/
fftn.rs

1use mlx_internal_macros::generate_macro;
2
3use crate::{
4    array::Array,
5    error::Result,
6    utils::{guard::Guarded, IntoOption},
7    with_stream, Stream,
8};
9
10use super::{
11    utils::{
12        legacy_fftn_options, resolve_lengths_and_axes, resolve_size_and_axis_unchecked,
13        resolve_sizes_and_axes_unchecked,
14    },
15    FftnOptions,
16};
17
18/// One dimensional discrete Fourier Transform.
19///
20/// # Params
21///
22/// - `a`: The input array.
23/// - `n`: Size of the transformed axis. The corresponding axis in the input is truncated or padded
24///   with zeros to match `n`. The default value is `a.shape[axis]`.
25/// - `axis`: Axis along which to perform the FFT. The default is -1.
26pub fn fft(
27    a: impl AsRef<Array>,
28    n: impl Into<Option<i32>>,
29    axis: impl Into<Option<i32>>,
30) -> Result<Array> {
31    let a = a.as_ref();
32    let (n, axis) = resolve_size_and_axis_unchecked(a, n.into(), axis.into());
33    let stream = Stream::thread_local_or_default();
34    Array::try_from_op(|res| unsafe {
35        mlx_sys::mlx_fft_fft(
36            res,
37            a.as_ptr(),
38            n,
39            axis,
40            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
41            stream.as_ref().as_ptr(),
42        )
43    })
44}
45
46/// Compatibility shim for [`fft`].
47#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
48#[deprecated(
49    since = "0.26.0",
50    note = "use `with_stream` or `with_device` around `fft`"
51)]
52pub fn fft_device(
53    a: impl AsRef<Array>,
54    #[optional] n: impl Into<Option<i32>>,
55    #[optional] axis: impl Into<Option<i32>>,
56    #[optional] stream: impl AsRef<Stream>,
57) -> Result<Array> {
58    with_stream(stream.as_ref(), || fft(a, n, axis))
59}
60
61/// Two dimensional discrete Fourier Transform.
62///
63/// # Params
64///
65/// - `a`: The input array.
66/// - `s`: Size of the transformed axes. The corresponding axes in the input are truncated or padded
67///   with zeros to match `s`. The default value is the sizes of `a` along `axes`.
68/// - `axes`: Axes along which to perform the FFT. The default is `[-2, -1]`.
69pub fn fft2<'a>(
70    a: impl AsRef<Array>,
71    s: impl IntoOption<&'a [i32]>,
72    axes: impl IntoOption<&'a [i32]>,
73) -> Result<Array> {
74    let a = a.as_ref();
75    let axes = axes.into_option().unwrap_or(&[-2, -1]);
76    let (s, axes) = resolve_sizes_and_axes_unchecked(a, s.into_option(), Some(axes));
77
78    let num_s = s.len();
79    let num_axes = axes.len();
80
81    let s_ptr = s.as_ptr();
82    let axes_ptr = axes.as_ptr();
83    let stream = Stream::thread_local_or_default();
84
85    Array::try_from_op(|res| unsafe {
86        mlx_sys::mlx_fft_fft2(
87            res,
88            a.as_ptr(),
89            s_ptr,
90            num_s,
91            axes_ptr,
92            num_axes,
93            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
94            stream.as_ref().as_ptr(),
95        )
96    })
97}
98
99/// Compatibility shim for [`fft2`].
100#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
101#[deprecated(
102    since = "0.26.0",
103    note = "use `with_stream` or `with_device` around `fft2`"
104)]
105pub fn fft2_device<'a>(
106    a: impl AsRef<Array>,
107    #[optional] s: impl IntoOption<&'a [i32]>,
108    #[optional] axes: impl IntoOption<&'a [i32]>,
109    #[optional] stream: impl AsRef<Stream>,
110) -> Result<Array> {
111    with_stream(stream.as_ref(), || fft2(a, s, axes))
112}
113
114/// n-dimensional discrete Fourier Transform.
115///
116/// # Params
117///
118/// - `a`: The input array.
119/// - `options`: Transform lengths and axes. The default transforms all axes at their input sizes.
120pub fn fftn(a: impl AsRef<Array>, options: FftnOptions) -> Result<Array> {
121    let a = a.as_ref();
122    let (s, axes) = resolve_lengths_and_axes(a.shape(), options.lengths.as_deref(), &options.axes)?;
123    let num_s = s.len();
124    let num_axes = axes.len();
125
126    let s_ptr = s.as_ptr();
127    let axes_ptr = axes.as_ptr();
128    let stream = Stream::thread_local_or_default();
129
130    Array::try_from_op(|res| unsafe {
131        mlx_sys::mlx_fft_fftn(
132            res,
133            a.as_ptr(),
134            s_ptr,
135            num_s,
136            axes_ptr,
137            num_axes,
138            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
139            stream.as_ref().as_ptr(),
140        )
141    })
142}
143
144/// Compatibility shim for [`fftn`].
145#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
146#[deprecated(
147    since = "0.26.0",
148    note = "use `with_stream` or `with_device` around `fftn` with `FftnOptions`"
149)]
150pub fn fftn_device<'a>(
151    a: impl AsRef<Array>,
152    #[optional] s: impl IntoOption<&'a [i32]>,
153    #[optional] axes: impl IntoOption<&'a [i32]>,
154    #[optional] stream: impl AsRef<Stream>,
155) -> Result<Array> {
156    let options = legacy_fftn_options(s.into_option(), axes.into_option())?;
157    with_stream(stream.as_ref(), || fftn(a, options))
158}
159
160/// One dimensional inverse discrete Fourier Transform.
161///
162/// # Params
163///
164/// - `a`: Input array.
165/// - `n`: Size of the transformed axis. The corresponding axis in the input is truncated or padded
166///   with zeros to match `n`. The default value is `a.shape[axis]` if not specified.
167/// - `axis`: Axis along which to perform the FFT. The default is `-1` if not specified.
168pub fn ifft(
169    a: impl AsRef<Array>,
170    n: impl Into<Option<i32>>,
171    axis: impl Into<Option<i32>>,
172) -> Result<Array> {
173    let a = a.as_ref();
174    let (n, axis) = resolve_size_and_axis_unchecked(a, n.into(), axis.into());
175    let stream = Stream::thread_local_or_default();
176
177    Array::try_from_op(|res| unsafe {
178        mlx_sys::mlx_fft_ifft(
179            res,
180            a.as_ptr(),
181            n,
182            axis,
183            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
184            stream.as_ref().as_ptr(),
185        )
186    })
187}
188
189/// Compatibility shim for [`ifft`].
190#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
191#[deprecated(
192    since = "0.26.0",
193    note = "use `with_stream` or `with_device` around `ifft`"
194)]
195pub fn ifft_device(
196    a: impl AsRef<Array>,
197    #[optional] n: impl Into<Option<i32>>,
198    #[optional] axis: impl Into<Option<i32>>,
199    #[optional] stream: impl AsRef<Stream>,
200) -> Result<Array> {
201    with_stream(stream.as_ref(), || ifft(a, n, axis))
202}
203
204/// Two dimensional inverse discrete Fourier Transform.
205///
206/// # Params
207///
208/// - `a`: The input array.
209/// - `s`: Size of the transformed axes. The corresponding axes in the input are truncated or padded
210///   with zeros to match `s`. The default value is the sizes of `a` along `axes`.
211/// - `axes`: Axes along which to perform the FFT. The default is `[-2, -1]`.
212pub fn ifft2<'a>(
213    a: impl AsRef<Array>,
214    s: impl IntoOption<&'a [i32]>,
215    axes: impl IntoOption<&'a [i32]>,
216) -> Result<Array> {
217    let a = a.as_ref();
218    let axes = axes.into_option().unwrap_or(&[-2, -1]);
219    let (s, axes) = resolve_sizes_and_axes_unchecked(a, s.into_option(), Some(axes));
220
221    let num_s = s.len();
222    let num_axes = axes.len();
223
224    let s_ptr = s.as_ptr();
225    let axes_ptr = axes.as_ptr();
226    let stream = Stream::thread_local_or_default();
227
228    Array::try_from_op(|res| unsafe {
229        mlx_sys::mlx_fft_ifft2(
230            res,
231            a.as_ptr(),
232            s_ptr,
233            num_s,
234            axes_ptr,
235            num_axes,
236            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
237            stream.as_ref().as_ptr(),
238        )
239    })
240}
241
242/// Compatibility shim for [`ifft2`].
243#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
244#[deprecated(
245    since = "0.26.0",
246    note = "use `with_stream` or `with_device` around `ifft2`"
247)]
248pub fn ifft2_device<'a>(
249    a: impl AsRef<Array>,
250    #[optional] s: impl IntoOption<&'a [i32]>,
251    #[optional] axes: impl IntoOption<&'a [i32]>,
252    #[optional] stream: impl AsRef<Stream>,
253) -> Result<Array> {
254    with_stream(stream.as_ref(), || ifft2(a, s, axes))
255}
256
257/// n-dimensional inverse discrete Fourier Transform.
258///
259/// # Params
260///
261/// - `a`: The input array.
262/// - `options`: Transform lengths and axes. The default transforms all axes at their input sizes.
263pub fn ifftn(a: impl AsRef<Array>, options: FftnOptions) -> Result<Array> {
264    let a = a.as_ref();
265    let (s, axes) = resolve_lengths_and_axes(a.shape(), options.lengths.as_deref(), &options.axes)?;
266    let num_s = s.len();
267    let num_axes = axes.len();
268
269    let s_ptr = s.as_ptr();
270    let axes_ptr = axes.as_ptr();
271    let stream = Stream::thread_local_or_default();
272
273    Array::try_from_op(|res| unsafe {
274        mlx_sys::mlx_fft_ifftn(
275            res,
276            a.as_ptr(),
277            s_ptr,
278            num_s,
279            axes_ptr,
280            num_axes,
281            mlx_sys::mlx_fft_norm__MLX_FFT_NORM_BACKWARD,
282            stream.as_ref().as_ptr(),
283        )
284    })
285}
286
287/// Compatibility shim for [`ifftn`].
288#[generate_macro(customize(forwarding_shim = true, root = "$crate::fft"))]
289#[deprecated(
290    since = "0.26.0",
291    note = "use `with_stream` or `with_device` around `ifftn` with `FftnOptions`"
292)]
293pub fn ifftn_device<'a>(
294    a: impl AsRef<Array>,
295    #[optional] s: impl IntoOption<&'a [i32]>,
296    #[optional] axes: impl IntoOption<&'a [i32]>,
297    #[optional] stream: impl AsRef<Stream>,
298) -> Result<Array> {
299    let options = legacy_fftn_options(s.into_option(), axes.into_option())?;
300    with_stream(stream.as_ref(), || ifftn(a, options))
301}
302
303#[cfg(test)]
304mod tests {
305    use crate::{
306        complex64, fft::*, ops::indexing::TryIndexOp, test_utils::assert_array_eq,
307        test_utils::tolerances, Array, Dtype, Stream,
308    };
309
310    #[test]
311    fn test_fft() {
312        const FFT_DATA: &[f32] = &[1.0, 2.0, 3.0, 4.0];
313        const FFT_SHAPE: &[i32] = &[4];
314        const FFT_EXPECTED: &[complex64; 4] = &[
315            complex64::new(10.0, 0.0),
316            complex64::new(-2.0, 2.0),
317            complex64::new(-2.0, 0.0),
318            complex64::new(-2.0, -2.0),
319        ];
320
321        let array = Array::from_slice(FFT_DATA, FFT_SHAPE);
322        let fft = fft(&array, None, None).unwrap();
323
324        assert_eq!(fft.dtype(), Dtype::Complex64);
325        assert_array_eq(
326            &fft,
327            Array::from_slice(FFT_EXPECTED, FFT_SHAPE),
328            tolerances::EXACT.rtol,
329            tolerances::EXACT.atol,
330        );
331
332        let ifft = ifft(&fft, None, None).unwrap();
333
334        assert_eq!(ifft.dtype(), Dtype::Complex64);
335        let expected = FFT_DATA
336            .iter()
337            .map(|&x| complex64::new(x, 0.0))
338            .collect::<Vec<_>>();
339        assert_array_eq(
340            &ifft,
341            Array::from_slice(&expected, FFT_SHAPE),
342            tolerances::EXACT.rtol,
343            tolerances::EXACT.atol,
344        );
345
346        assert_array_eq(
347            array,
348            Array::from_slice(FFT_DATA, FFT_SHAPE),
349            tolerances::EXACT.rtol,
350            tolerances::EXACT.atol,
351        );
352    }
353
354    #[test]
355    fn test_fft2() {
356        const FFT2_DATA: &[f32] = &[1.0, 1.0, 1.0, 1.0];
357        const FFT2_SHAPE: &[i32] = &[2, 2];
358        const FFT2_EXPECTED: &[complex64; 4] = &[
359            complex64::new(4.0, 0.0),
360            complex64::new(0.0, 0.0),
361            complex64::new(0.0, 0.0),
362            complex64::new(0.0, 0.0),
363        ];
364
365        let array = Array::from_slice(FFT2_DATA, FFT2_SHAPE);
366        let fft2 = fft2(&array, None, None).unwrap();
367
368        assert_eq!(fft2.dtype(), Dtype::Complex64);
369        assert_array_eq(
370            &fft2,
371            Array::from_slice(FFT2_EXPECTED, FFT2_SHAPE),
372            tolerances::EXACT.rtol,
373            tolerances::EXACT.atol,
374        );
375
376        let ifft2 = ifft2(&fft2, None, None).unwrap();
377
378        assert_eq!(ifft2.dtype(), Dtype::Complex64);
379        let expected = FFT2_DATA
380            .iter()
381            .map(|&x| complex64::new(x, 0.0))
382            .collect::<Vec<_>>();
383        assert_array_eq(
384            &ifft2,
385            Array::from_slice(&expected, FFT2_SHAPE),
386            tolerances::EXACT.rtol,
387            tolerances::EXACT.atol,
388        );
389
390        assert_array_eq(
391            array,
392            Array::from_slice(FFT2_DATA, FFT2_SHAPE),
393            tolerances::EXACT.rtol,
394            tolerances::EXACT.atol,
395        );
396    }
397
398    #[test]
399    fn test_fftn() {
400        const FFTN_DATA: &[f32] = &[1.0; 8];
401        const FFTN_SHAPE: &[i32] = &[2, 2, 2];
402        const FFTN_EXPECTED: &[complex64; 8] = &[
403            complex64::new(8.0, 0.0),
404            complex64::new(0.0, 0.0),
405            complex64::new(0.0, 0.0),
406            complex64::new(0.0, 0.0),
407            complex64::new(0.0, 0.0),
408            complex64::new(0.0, 0.0),
409            complex64::new(0.0, 0.0),
410            complex64::new(0.0, 0.0),
411        ];
412
413        let array = Array::from_slice(FFTN_DATA, FFTN_SHAPE);
414        let fftn = fftn(&array, FftnOptions::default()).unwrap();
415
416        assert_eq!(fftn.dtype(), Dtype::Complex64);
417        assert_array_eq(
418            &fftn,
419            Array::from_slice(FFTN_EXPECTED, FFTN_SHAPE),
420            tolerances::EXACT.rtol,
421            tolerances::EXACT.atol,
422        );
423
424        let ifftn = ifftn(
425            &fftn,
426            FftnOptions {
427                lengths: Some(FFTN_SHAPE.to_vec()),
428                axes: [0, 1, 2].into(),
429            },
430        )
431        .unwrap();
432
433        assert_eq!(ifftn.dtype(), Dtype::Complex64);
434        let expected = FFTN_DATA
435            .iter()
436            .map(|&x| complex64::new(x, 0.0))
437            .collect::<Vec<_>>();
438        assert_array_eq(
439            &ifftn,
440            Array::from_slice(&expected, FFTN_SHAPE),
441            tolerances::EXACT.rtol,
442            tolerances::EXACT.atol,
443        );
444
445        assert_array_eq(
446            array,
447            Array::from_slice(FFTN_DATA, FFTN_SHAPE),
448            tolerances::EXACT.rtol,
449            tolerances::EXACT.atol,
450        );
451    }
452
453    #[test]
454    fn asymmetric_multidimensional_transforms_use_backward_normalization() {
455        let input2 = Array::from_slice(&[1.0_f32, 2.0, 3.0, 5.0, 7.0, 11.0], &[2, 3]);
456        let spectrum2 = fft2(&input2, None, None).unwrap();
457        assert_array_eq(
458            spectrum2.try_index((0, 0)).unwrap(),
459            Array::from(complex64::new(29.0, 0.0)),
460            tolerances::EXACT.rtol,
461            tolerances::EXACT.atol,
462        );
463        // The inverse of a complex spectrum is complex with ~zero imaginaries;
464        // compare in the output's dtype rather than casting it away.
465        let roundtrip2 = ifft2(&spectrum2, None, None).unwrap();
466        let expected2 = Array::from_slice(
467            &[1.0_f32, 2.0, 3.0, 5.0, 7.0, 11.0].map(|re| complex64::new(re, 0.0)),
468            &[2, 3],
469        );
470        assert_array_eq(
471            &roundtrip2,
472            &expected2,
473            tolerances::STANDARD.rtol,
474            tolerances::STANDARD.atol,
475        );
476
477        let inputn = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 2, 2]);
478        let spectrumn = fftn(&inputn, FftnOptions::default()).unwrap();
479        assert_array_eq(
480            spectrumn.try_index((0, 0, 0)).unwrap(),
481            Array::from(complex64::new(36.0, 0.0)),
482            tolerances::EXACT.rtol,
483            tolerances::EXACT.atol,
484        );
485        let roundtripn = ifftn(&spectrumn, FftnOptions::default()).unwrap();
486        let expectedn = Array::from_slice(
487            &[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0].map(|re| complex64::new(re, 0.0)),
488            &[2, 2, 2],
489        );
490        assert_array_eq(
491            &roundtripn,
492            &expectedn,
493            tolerances::STANDARD.rtol,
494            tolerances::STANDARD.atol,
495        );
496    }
497
498    #[allow(deprecated)]
499    #[test]
500    fn compatibility_functions_and_macros_forward_to_canonical_fft() {
501        let input = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
502        let stream = Stream::cpu();
503        let canonical = fft(&input, None, None).unwrap();
504
505        for compatibility in [
506            fft_device(&input, None, None, &stream).unwrap(),
507            fft!(&input).unwrap(),
508            fft!(&input, stream = &stream).unwrap(),
509        ] {
510            assert_array_eq(
511                &compatibility,
512                &canonical,
513                tolerances::EXACT.rtol,
514                tolerances::EXACT.atol,
515            );
516        }
517
518        let canonical_n = fftn(
519            &input,
520            FftnOptions {
521                lengths: Some(vec![4]),
522                axes: (-1).into(),
523            },
524        )
525        .unwrap();
526        let compatibility_n = fftn!(&input, s = &[4]).unwrap();
527        assert_array_eq(
528            compatibility_n,
529            canonical_n,
530            tolerances::EXACT.rtol,
531            tolerances::EXACT.atol,
532        );
533    }
534}
mlx_rs/fft/fftn.rs

mlx_rs/fft/
fftn.rs
