zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-06 20:20:11 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Switch to nanobind (#839)

Browse Source 
* mostly builds

* most tests pass

* fix circle build

* add back buffer protocol

* includes

* fix for py38

* limit to cpu device

* include

* fix stubs

* move signatures for docs

* stubgen + docs fix

* doc for compiled function, comments
This commit is contained in:
Awni Hannun
2024-03-18 20:12:25 -07:00
committed by GitHub
parent d39ed54f8e
commit 9a8ee00246
34 changed files with 2343 additions and 2344 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
python/src/fft.cpp
View File

@@ -1,19 +1,20 @@
// Copyright © 2023 Apple Inc.
// Copyright © 2023-2024 Apple Inc.
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include "python/src/utils.h"
#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"
namespace py = pybind11;
using namespace py::literals;
namespace nb = nanobind;
using namespace nb::literals;
using namespace mlx::core;
void init_fft(py::module_& parent_module) {
void init_fft(nb::module_& parent_module) {
auto m = parent_module.def_submodule(
"fft", "mlx.core.fft: Fast Fourier Transforms.");
m.def(
@@ -29,9 +30,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"n"_a = none,
"n"_a = nb::none(),
"axis"_a = -1,
"stream"_a = none,
"stream"_a = nb::none(),
R"pbdoc(
One dimensional discrete Fourier Transform.
@@ -59,9 +60,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"n"_a = none,
"n"_a = nb::none(),
"axis"_a = -1,
"stream"_a = none,
"stream"_a = nb::none(),
R"pbdoc(
One dimensional inverse discrete Fourier Transform.
@@ -95,9 +96,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = std::vector<int>{-2, -1},
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a.none() = std::vector<int>{-2, -1},
"stream"_a = nb::none(),
R"pbdoc(
Two dimensional discrete Fourier Transform.
@@ -132,9 +133,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = std::vector<int>{-2, -1},
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a.none() = std::vector<int>{-2, -1},
"stream"_a = nb::none(),
R"pbdoc(
Two dimensional inverse discrete Fourier Transform.
@@ -169,9 +170,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = none,
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
n-dimensional discrete Fourier Transform.
@@ -207,9 +208,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = none,
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
n-dimensional inverse discrete Fourier Transform.
@@ -239,9 +240,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"n"_a = none,
"n"_a = nb::none(),
"axis"_a = -1,
"stream"_a = none,
"stream"_a = nb::none(),
R"pbdoc(
One dimensional discrete Fourier Transform on a real input.
@@ -274,9 +275,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"n"_a = none,
"n"_a = nb::none(),
"axis"_a = -1,
"stream"_a = none,
"stream"_a = nb::none(),
R"pbdoc(
The inverse of :func:`rfft`.
@@ -314,9 +315,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = std::vector<int>{-2, -1},
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a.none() = std::vector<int>{-2, -1},
"stream"_a = nb::none(),
R"pbdoc(
Two dimensional real discrete Fourier Transform.
@@ -357,9 +358,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = std::vector<int>{-2, -1},
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a.none() = std::vector<int>{-2, -1},
"stream"_a = nb::none(),
R"pbdoc(
The inverse of :func:`rfft2`.
@@ -400,9 +401,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = none,
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
n-dimensional real discrete Fourier Transform.
@@ -443,9 +444,9 @@ void init_fft(py::module_& parent_module) {
}
},
"a"_a,
"s"_a = none,
"axes"_a = none,
"stream"_a = none,
"s"_a = nb::none(),
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
The inverse of :func:`rfftn`.