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

Add SVD primitive (#809)

Browse Source 
Add SVD op using Accelerate's LAPACK following
https://developer.apple.com/documentation/accelerate/
compressing_an_image_using_linear_algebra

Co-authored-by: Nicolo Valigi <nvaligi@apple.com>
This commit is contained in:
nicolov
2024-03-12 20:30:11 +01:00
committed by GitHub
parent ffb19df3c0
commit d0c544a868
13 changed files with 318 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/src/linalg.cpp
View File

@@ -16,6 +16,13 @@ using namespace py::literals;
using namespace mlx::core;
using namespace mlx::core::linalg;
namespace {
py::tuple svd_helper(const array& a, StreamOrDevice s /* = {} */) {
const auto result = svd(a, s);
return py::make_tuple(result.at(0), result.at(1), result.at(2));
}
} // namespace
void init_linalg(py::module_& parent_module) {
py::options options;
options.disable_function_signatures();
@@ -186,7 +193,7 @@ void init_linalg(py::module_& parent_module) {
R"pbdoc(
qr(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array)
The QR factorizatoin of the input matrix.
The QR factorization of the input matrix.
This function supports arrays with at least 2 dimensions. The matrices
which are factorized are assumed to be in the last two dimensions of
@@ -210,4 +217,28 @@ void init_linalg(py::module_& parent_module) {
array([[-2.23607, -3.57771],
[0, 0.447214]], dtype=float32)
)pbdoc");
m.def(
"svd",
&svd_helper,
"a"_a,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
svd(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array, array)
The Singular Value Decomposition (SVD) of the input matrix.
This function supports arrays with at least 2 dimensions. When the input
has more than two dimensions, the function iterates over all indices of the first
a.ndim - 2 dimensions and for each combination SVD is applied to the last two indices.
Args:
a (array): Input array.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
tuple(array, array, array): The ``U``, ``S``, and ``Vt`` matrices, such that
``A = U @ diag(S) @ Vt``
)pbdoc");
}