Module fft

Expand description

Fast Fourier Transform (FFT) and its inverse (IFFT) for one, two, and N dimensions.

Unsuffixed functions are the canonical named Result paths and use the current thread’s scoped stream or runtime default. Select an explicit stream or device with crate::with_stream or crate::with_device. The _device functions and generated operation macros are deprecated forwarding shims.

N-dimensional transforms accept FftnOptions. Its default selects every axis at the input lengths, and correlated lengths and axes are validated before calling MLX.

§Examples

§One dimension

use mlx_rs::{
    Dtype, Array, complex64, fft::*,
    test_utils::{assert_array_eq, tolerances},
};

let src = [1.0f32, 2.0, 3.0, 4.0];
let mut array = Array::from_slice(&src[..], &[4]);

let mut fft_result = fft(&array, 4, 0).unwrap();
assert_eq!(fft_result.dtype(), Dtype::Complex64);

let expected = Array::from_slice(&[
    complex64::new(10.0, 0.0),
    complex64::new(-2.0, 2.0),
    complex64::new(-2.0, 0.0),
    complex64::new(-2.0, -2.0),
], &[4]);
assert_array_eq(&fft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut ifft_result = ifft(&fft_result, 4, 0).unwrap();
assert_eq!(ifft_result.dtype(), Dtype::Complex64);

let expected = Array::from_slice(&[
   complex64::new(1.0, 0.0),
   complex64::new(2.0, 0.0),
   complex64::new(3.0, 0.0),
   complex64::new(4.0, 0.0),
], &[4]);
assert_array_eq(ifft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut rfft_result = rfft(&array, 4, 0).unwrap();
assert_eq!(rfft_result.dtype(), Dtype::Complex64);

let expected = Array::from_slice(&[
   complex64::new(10.0, 0.0),
   complex64::new(-2.0, 2.0),
   complex64::new(-2.0, 0.0),
], &[3]);
assert_array_eq(&rfft_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut irfft_result = irfft(&rfft_result, 4, 0).unwrap();
assert_eq!(irfft_result.dtype(), Dtype::Float32);
assert_array_eq(
    irfft_result,
    Array::from_slice(&src, &[4]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

// The original array is not modified
assert_array_eq(
    array,
    Array::from_slice(&src, &[4]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

§Two dimensions

use mlx_rs::{
    Dtype, Array, complex64, fft::*,
    test_utils::{assert_array_eq, tolerances},
};

let src = [1.0f32, 1.0, 1.0, 1.0];
let mut array = Array::from_slice(&src[..], &[2, 2]);

let mut fft2_result = fft2(&array, None, None).unwrap();
assert_eq!(fft2_result.dtype(), Dtype::Complex64);
let expected = Array::from_slice(&[
    complex64::new(4.0, 0.0),
    complex64::new(0.0, 0.0),
    complex64::new(0.0, 0.0),
    complex64::new(0.0, 0.0),
], &[2, 2]);
assert_array_eq(&fft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut ifft2_result = ifft2(&fft2_result, None, None).unwrap();
assert_eq!(ifft2_result.dtype(), Dtype::Complex64);

let expected = Array::from_slice(&[
   complex64::new(1.0, 0.0),
   complex64::new(1.0, 0.0),
   complex64::new(1.0, 0.0),
   complex64::new(1.0, 0.0),
], &[2, 2]);
assert_array_eq(ifft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut rfft2_result = rfft2(&array, None, None).unwrap();
assert_eq!(rfft2_result.dtype(), Dtype::Complex64);

let expected = Array::from_slice(&[
    complex64::new(4.0, 0.0),
    complex64::new(0.0, 0.0),
    complex64::new(0.0, 0.0),
    complex64::new(0.0, 0.0),
], &[2, 2]);
assert_array_eq(&rfft2_result, expected, tolerances::EXACT.rtol, tolerances::EXACT.atol);

let mut irfft2_result = irfft2(&rfft2_result, None, None).unwrap();
assert_eq!(irfft2_result.dtype(), Dtype::Float32);
assert_array_eq(
    irfft2_result,
    Array::from_slice(&src, &[2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

// The original array is not modified
assert_array_eq(
    array,
    Array::from_slice(&src, &[2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

§`N` dimensions

use mlx_rs::{
    Dtype, Array, complex64, fft::*,
    test_utils::{assert_array_eq, tolerances},
};

let mut array = Array::ones::<f32>(&[2, 2, 2]).unwrap();
let mut fftn_result = fftn(&array, FftnOptions::default()).unwrap();
assert_eq!(fftn_result.dtype(), Dtype::Complex64);

let mut expected = [complex64::new(0.0, 0.0); 8];
expected[0] = complex64::new(8.0, 0.0);
assert_array_eq(
    &fftn_result,
    Array::from_slice(&expected, &[2, 2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

let mut ifftn_result = ifftn(&fftn_result, FftnOptions::default()).unwrap();
assert_eq!(ifftn_result.dtype(), Dtype::Complex64);

let expected = [complex64::new(1.0, 0.0); 8];
assert_array_eq(
    ifftn_result,
    Array::from_slice(&expected, &[2, 2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

let mut rfftn_result = rfftn(&array, FftnOptions::default()).unwrap();
assert_eq!(rfftn_result.dtype(), Dtype::Complex64);

let mut expected = [complex64::new(0.0, 0.0); 8];
expected[0] = complex64::new(8.0, 0.0);
assert_array_eq(
    &rfftn_result,
    Array::from_slice(&expected, &[2, 2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

let mut irfftn_result = irfftn(&rfftn_result, FftnOptions::default()).unwrap();
assert_eq!(irfftn_result.dtype(), Dtype::Float32);

let expected = [1.0; 8];
assert_array_eq(
    irfftn_result,
    Array::from_slice(&expected, &[2, 2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

// The original array is not modified
assert_array_eq(
    array,
    Array::from_slice(&[1.0; 8], &[2, 2, 2]),
    tolerances::EXACT.rtol,
    tolerances::EXACT.atol,
);

Macros§

fftDeprecated: Macro generated for the function crate::fft::fft. See the function documentation for more details.
fft2Deprecated: Macro generated for the function crate::fft::fft2. See the function documentation for more details.
fftnDeprecated: Macro generated for the function crate::fft::fftn. See the function documentation for more details.
fftshiftDeprecated: Macro generated for the function crate::fft::fftshift. See the function documentation for more details.
ifftDeprecated: Macro generated for the function crate::fft::ifft. See the function documentation for more details.
ifft2Deprecated: Macro generated for the function crate::fft::ifft2. See the function documentation for more details.
ifftnDeprecated: Macro generated for the function crate::fft::ifftn. See the function documentation for more details.
ifftshiftDeprecated: Macro generated for the function crate::fft::ifftshift. See the function documentation for more details.
irfftDeprecated: Macro generated for the function crate::fft::irfft. See the function documentation for more details.
irfft2Deprecated: Macro generated for the function crate::fft::irfft2. See the function documentation for more details.
irfftnDeprecated: Macro generated for the function crate::fft::irfftn. See the function documentation for more details.
rfftDeprecated: Macro generated for the function crate::fft::rfft. See the function documentation for more details.
rfft2Deprecated: Macro generated for the function crate::fft::rfft2. See the function documentation for more details.
rfftnDeprecated: Macro generated for the function crate::fft::rfftn. See the function documentation for more details.

Structs§

FftnOptions: Options shared by the n-dimensional FFT and real FFT transforms.

Functions§

fft: One dimensional discrete Fourier Transform.
fft2: Two dimensional discrete Fourier Transform.
fft2_deviceDeprecated: Compatibility shim for fft2.
fft_deviceDeprecated: Compatibility shim for fft.
fftfreq: Returns the discrete Fourier transform sample frequencies for a transform of length n.
fftn: n-dimensional discrete Fourier Transform.
fftn_deviceDeprecated: Compatibility shim for fftn.
fftshift: Shift the zero-frequency component to the center of the spectrum.
fftshift_deviceDeprecated: Compatibility shim for fftshift.
ifft: One dimensional inverse discrete Fourier Transform.
ifft2: Two dimensional inverse discrete Fourier Transform.
ifft2_deviceDeprecated: Compatibility shim for ifft2.
ifft_deviceDeprecated: Compatibility shim for ifft.
ifftn: n-dimensional inverse discrete Fourier Transform.
ifftn_deviceDeprecated: Compatibility shim for ifftn.
ifftshift: The inverse of fftshift.
ifftshift_deviceDeprecated: Compatibility shim for ifftshift.
irfft: The inverse of rfft().
irfft2: The inverse of rfft2().
irfft2_deviceDeprecated: Compatibility shim for irfft2.
irfft_deviceDeprecated: Compatibility shim for irfft.
irfftn: The inverse of rfftn().
irfftn_deviceDeprecated: Compatibility shim for irfftn.
rfft: One dimensional discrete Fourier Transform on a real input.
rfft2: Two-dimensional real discrete Fourier Transform.
rfft2_deviceDeprecated: Compatibility shim for rfft2.
rfft_deviceDeprecated: Compatibility shim for rfft.
rfftfreq: Returns the nonnegative discrete Fourier transform sample frequencies for a real transform.
rfftn: n-dimensional real discrete Fourier Transform.
rfftn_deviceDeprecated: Compatibility shim for rfftn.

Module fft

Module fft Copy item path

§Examples

§One dimension

§Two dimensions

§N dimensions

Macros§

Structs§

Functions§

Module fft

§`N` dimensions