| | |
| |
|
| | #include <nanobind/nanobind.h> |
| | #include <nanobind/stl/optional.h> |
| | #include <nanobind/stl/variant.h> |
| | #include <nanobind/stl/vector.h> |
| | #include <numeric> |
| |
|
| | #include "mlx/fft.h" |
| | #include "mlx/ops.h" |
| | #include "python/src/small_vector.h" |
| |
|
| | namespace mx = mlx::core; |
| | namespace nb = nanobind; |
| | using namespace nb::literals; |
| |
|
| | void init_fft(nb::module_& parent_module) { |
| | auto m = parent_module.def_submodule( |
| | "fft", "mlx.core.fft: Fast Fourier Transforms."); |
| | m.def( |
| | "fft", |
| | [](const mx::array& a, |
| | const std::optional<int>& n, |
| | int axis, |
| | mx::StreamOrDevice s) { |
| | if (n.has_value()) { |
| | return mx::fft::fft(a, n.value(), axis, s); |
| | } else { |
| | return mx::fft::fft(a, axis, s); |
| | } |
| | }, |
| | "a"_a, |
| | "n"_a = nb::none(), |
| | "axis"_a = -1, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | One dimensional discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | n (int, optional): Size of the transformed axis. The |
| | corresponding axis in the input is truncated or padded with |
| | zeros to match ``n``. The default value is ``a.shape[axis]``. |
| | axis (int, optional): Axis along which to perform the FFT. The |
| | default is ``-1``. |
| | |
| | Returns: |
| | array: The DFT of the input along the given axis. |
| | )pbdoc"); |
| | m.def( |
| | "ifft", |
| | [](const mx::array& a, |
| | const std::optional<int>& n, |
| | int axis, |
| | mx::StreamOrDevice s) { |
| | if (n.has_value()) { |
| | return mx::fft::ifft(a, n.value(), axis, s); |
| | } else { |
| | return mx::fft::ifft(a, axis, s); |
| | } |
| | }, |
| | "a"_a, |
| | "n"_a = nb::none(), |
| | "axis"_a = -1, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | One dimensional inverse discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | n (int, optional): Size of the transformed axis. The |
| | corresponding axis in the input is truncated or padded with |
| | zeros to match ``n``. The default value is ``a.shape[axis]``. |
| | axis (int, optional): Axis along which to perform the FFT. The |
| | default is ``-1``. |
| | |
| | Returns: |
| | array: The inverse DFT of the input along the given axis. |
| | )pbdoc"); |
| | m.def( |
| | "fft2", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::fftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::fftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[fft2] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::fftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a.none() = std::vector<int>{-2, -1}, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | Two dimensional discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``[-2, -1]``. |
| | |
| | Returns: |
| | array: The DFT of the input along the given axes. |
| | )pbdoc"); |
| | m.def( |
| | "ifft2", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::ifftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::ifftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[ifft2] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::ifftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a.none() = std::vector<int>{-2, -1}, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | Two dimensional inverse discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``[-2, -1]``. |
| | |
| | Returns: |
| | array: The inverse DFT of the input along the given axes. |
| | )pbdoc"); |
| | m.def( |
| | "fftn", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::fftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::fftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[fftn] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::fftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | n-dimensional discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``None`` in which case the FFT is over the last |
| | ``len(s)`` axes are or all axes if ``s`` is also ``None``. |
| | |
| | Returns: |
| | array: The DFT of the input along the given axes. |
| | )pbdoc"); |
| | m.def( |
| | "ifftn", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::ifftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::ifftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[ifftn] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::ifftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | n-dimensional inverse discrete Fourier Transform. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``None`` in which case the FFT is over the last |
| | ``len(s)`` axes or all axes if ``s`` is also ``None``. |
| | |
| | Returns: |
| | array: The inverse DFT of the input along the given axes. |
| | )pbdoc"); |
| | m.def( |
| | "rfft", |
| | [](const mx::array& a, |
| | const std::optional<int>& n, |
| | int axis, |
| | mx::StreamOrDevice s) { |
| | if (n.has_value()) { |
| | return mx::fft::rfft(a, n.value(), axis, s); |
| | } else { |
| | return mx::fft::rfft(a, axis, s); |
| | } |
| | }, |
| | "a"_a, |
| | "n"_a = nb::none(), |
| | "axis"_a = -1, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | One dimensional discrete Fourier Transform on a real input. |
| | |
| | The output has the same shape as the input except along ``axis`` in |
| | which case it has size ``n // 2 + 1``. |
| | |
| | Args: |
| | a (array): The input array. If the array is complex it will be silently |
| | cast to a real type. |
| | n (int, optional): Size of the transformed axis. The |
| | corresponding axis in the input is truncated or padded with |
| | zeros to match ``n``. The default value is ``a.shape[axis]``. |
| | axis (int, optional): Axis along which to perform the FFT. The |
| | default is ``-1``. |
| | |
| | Returns: |
| | array: The DFT of the input along the given axis. The output |
| | data type will be complex. |
| | )pbdoc"); |
| | m.def( |
| | "irfft", |
| | [](const mx::array& a, |
| | const std::optional<int>& n, |
| | int axis, |
| | mx::StreamOrDevice s) { |
| | if (n.has_value()) { |
| | return mx::fft::irfft(a, n.value(), axis, s); |
| | } else { |
| | return mx::fft::irfft(a, axis, s); |
| | } |
| | }, |
| | "a"_a, |
| | "n"_a = nb::none(), |
| | "axis"_a = -1, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | The inverse of :func:`rfft`. |
| | |
| | The output has the same shape as the input except along ``axis`` in |
| | which case it has size ``n``. |
| | |
| | Args: |
| | a (array): The input array. |
| | n (int, optional): Size of the transformed axis. The |
| | corresponding axis in the input is truncated or padded with |
| | zeros to match ``n // 2 + 1``. The default value is |
| | ``a.shape[axis] // 2 + 1``. |
| | axis (int, optional): Axis along which to perform the FFT. The |
| | default is ``-1``. |
| | |
| | Returns: |
| | array: The real array containing the inverse of :func:`rfft`. |
| | )pbdoc"); |
| | m.def( |
| | "rfft2", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::rfftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::rfftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[rfft2] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::rfftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a.none() = std::vector<int>{-2, -1}, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | Two dimensional real discrete Fourier Transform. |
| | |
| | The output has the same shape as the input except along the dimensions in |
| | ``axes`` in which case it has sizes from ``s``. The last axis in ``axes`` is |
| | treated as the real axis and will have size ``s[-1] // 2 + 1``. |
| | |
| | Args: |
| | a (array): The input array. If the array is complex it will be silently |
| | cast to a real type. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``[-2, -1]``. |
| | |
| | Returns: |
| | array: The real DFT of the input along the given axes. The output |
| | data type will be complex. |
| | )pbdoc"); |
| | m.def( |
| | "irfft2", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::irfftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::irfftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[irfft2] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::irfftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a.none() = std::vector<int>{-2, -1}, |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | The inverse of :func:`rfft2`. |
| | |
| | Note the input is generally complex. The dimensions of the input |
| | specified in ``axes`` are padded or truncated to match the sizes |
| | from ``s``. The last axis in ``axes`` is treated as the real axis |
| | and will have size ``s[-1] // 2 + 1``. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s`` except for the last axis |
| | which has size ``s[-1] // 2 + 1``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``[-2, -1]``. |
| | |
| | Returns: |
| | array: The real array containing the inverse of :func:`rfft2`. |
| | )pbdoc"); |
| | m.def( |
| | "rfftn", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::rfftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::rfftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[rfftn] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::rfftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | n-dimensional real discrete Fourier Transform. |
| | |
| | The output has the same shape as the input except along the dimensions in |
| | ``axes`` in which case it has sizes from ``s``. The last axis in ``axes`` is |
| | treated as the real axis and will have size ``s[-1] // 2 + 1``. |
| | |
| | Args: |
| | a (array): The input array. If the array is complex it will be silently |
| | cast to a real type. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``None`` in which case the FFT is over the last |
| | ``len(s)`` axes or all axes if ``s`` is also ``None``. |
| | |
| | Returns: |
| | array: The real DFT of the input along the given axes. The output |
| | )pbdoc"); |
| | m.def( |
| | "irfftn", |
| | [](const mx::array& a, |
| | const std::optional<mx::Shape>& n, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value() && n.has_value()) { |
| | return mx::fft::irfftn(a, n.value(), axes.value(), s); |
| | } else if (axes.has_value()) { |
| | return mx::fft::irfftn(a, axes.value(), s); |
| | } else if (n.has_value()) { |
| | throw std::invalid_argument( |
| | "[irfftn] `axes` should not be `None` if `s` is not `None`."); |
| | } else { |
| | return mx::fft::irfftn(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "s"_a = nb::none(), |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | The inverse of :func:`rfftn`. |
| | |
| | Note the input is generally complex. The dimensions of the input |
| | specified in ``axes`` are padded or truncated to match the sizes |
| | from ``s``. The last axis in ``axes`` is treated as the real axis |
| | and will have size ``s[-1] // 2 + 1``. |
| | |
| | Args: |
| | a (array): The input array. |
| | s (list(int), optional): Sizes of the transformed axes. The |
| | corresponding axes in the input are truncated or padded with |
| | zeros to match the sizes in ``s``. The default value is the |
| | sizes of ``a`` along ``axes``. |
| | axes (list(int), optional): Axes along which to perform the FFT. |
| | The default is ``None`` in which case the FFT is over the last |
| | ``len(s)`` axes or all axes if ``s`` is also ``None``. |
| | |
| | Returns: |
| | array: The real array containing the inverse of :func:`rfftn`. |
| | )pbdoc"); |
| | m.def( |
| | "fftshift", |
| | [](const mx::array& a, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value()) { |
| | return mx::fft::fftshift(a, axes.value(), s); |
| | } else { |
| | return mx::fft::fftshift(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | Shift the zero-frequency component to the center of the spectrum. |
| | |
| | Args: |
| | a (array): The input array. |
| | axes (list(int), optional): Axes over which to perform the shift. |
| | If ``None``, shift all axes. |
| | |
| | Returns: |
| | array: The shifted array with the same shape as the input. |
| | )pbdoc"); |
| | m.def( |
| | "ifftshift", |
| | [](const mx::array& a, |
| | const std::optional<std::vector<int>>& axes, |
| | mx::StreamOrDevice s) { |
| | if (axes.has_value()) { |
| | return mx::fft::ifftshift(a, axes.value(), s); |
| | } else { |
| | return mx::fft::ifftshift(a, s); |
| | } |
| | }, |
| | "a"_a, |
| | "axes"_a = nb::none(), |
| | "stream"_a = nb::none(), |
| | R"pbdoc( |
| | The inverse of :func:`fftshift`. While identical to :func:`fftshift` for even-length axes, |
| | the behavior differs for odd-length axes. |
| | |
| | Args: |
| | a (array): The input array. |
| | axes (list(int), optional): Axes over which to perform the inverse shift. |
| | If ``None``, shift all axes. |
| | |
| | Returns: |
| | array: The inverse-shifted array with the same shape as the input. |
| | )pbdoc"); |
| | } |
| |
|
