zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +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

30
mlx/linalg.cpp
View File

@@ -4,8 +4,9 @@
#include <ostream>
#include <vector>
#include "mlx/dtype.h"
#include "mlx/linalg.h"
#include "mlx/primitives.h"
#include "mlx/utils.h"
namespace mlx::core::linalg {
@@ -172,4 +173,31 @@ array norm(
return matrix_norm(a, ord, ax, keepdims, s);
}
std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
if (a.dtype() != float32) {
std::ostringstream msg;
msg << "[linalg::qr] Arrays must type float32. Received array "
<< "with type " << a.dtype() << ".";
throw std::invalid_argument(msg.str());
}
if (a.ndim() < 2) {
std::ostringstream msg;
msg << "[linalg::qr] Arrays must have >= 2 dimensions. Received array "
"with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
if (a.shape(-1) != a.shape(-2)) {
throw std::invalid_argument(
"[linalg::qr] Support for non-square matrices NYI.");
}
auto out = array::make_arrays(
{a.shape(), a.shape()},
{a.dtype(), a.dtype()},
std::make_unique<QRF>(to_stream(s)),
{astype(a, a.dtype(), s)});
return std::make_pair(out[0], out[1]);
}
} // namespace mlx::core::linalg