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

QR factorization (#310)

Browse Source 
* add qr factorization

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
taher
2024-01-26 09:27:31 -08:00
committed by GitHub
parent 2463496471
commit 077c1ee64a
20 changed files with 322 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
python/src/linalg.cpp
View File

@@ -177,4 +177,37 @@ void init_linalg(py::module_& parent_module) {
>>> la.norm(m[0, :, :]), LA.norm(m[1, :, :])
(array(3.74166, dtype=float32), array(11.225, dtype=float32))
)pbdoc");
m.def(
"qr",
&qr,
"a"_a,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
qr(a: array, *, stream: Union[None, Stream, Device] = None) -> (array, array)
The QR factorizatoin 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
the input.
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): The ``Q`` and ``R`` matrices.
Example:
>>> A = mx.array([[2., 3.], [1., 2.]])
>>> Q, R = mx.linalg.qr(A, stream=mx.cpu)
>>> Q
array([[-0.894427, -0.447214],
[-0.447214, 0.894427]], dtype=float32)
>>> R
array([[-2.23607, -3.57771],
[0, 0.447214]], dtype=float32)
)pbdoc");
}

4
python/src/ops.cpp
View File

@@ -55,7 +55,7 @@ void init_ops(py::module_& m) {
Args:
a (array): Input array.
shape (tuple(int)): New shape.
stream (Stream, optional): Stream or device. Defaults to ```None```
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
@@ -112,7 +112,7 @@ void init_ops(py::module_& m) {
Args:
a (array): Input array.
axis (int or tuple(int), optional): Axes to remove. Defaults
to ```None``` in which case all size one axes are removed.
to ``None`` in which case all size one axes are removed.
Returns:
array: The output array with size one axes removed.