mlx 1.0.0

data/mlx/python/src/fft.cpp

@@ -0,0 +1,514 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#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");
+}