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

non square qr (#1783)

Browse Source
This commit is contained in:
Awni Hannun
2025-01-21 14:07:47 -08:00
committed by GitHub
parent 1f4c127fb9
commit e6a7ab9675
3 changed files with 37 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
mlx/linalg.cpp
View File

@@ -187,13 +187,13 @@ std::pair<array, array> qr(const array& a, StreamOrDevice s /* = {} */) {
<< 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.");
}
int k = std::min(a.shape(-2), a.shape(-1));
auto q_shape = a.shape();
q_shape.back() = k;
auto r_shape = a.shape();
r_shape[r_shape.size() - 2] = k;
auto out = array::make_arrays(
{a.shape(), a.shape()},
{std::move(q_shape), std::move(r_shape)},
{a.dtype(), a.dtype()},
std::make_shared<QRF>(to_stream(s)),
{astype(a, a.dtype(), s)});