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

CPU mx.linalg.cholesky_inverse and mx.linalg.tri_inv (#1307)

Browse Source 
* add cholesky inv + tri inv

* always run tri_inv on cpu

* consistent naming
This commit is contained in:
Alex Barron
2024-08-08 15:18:02 -07:00
committed by GitHub
parent 780c197f95
commit 32668a7317
7 changed files with 267 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
mlx/linalg.cpp
View File

@@ -238,7 +238,7 @@ std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
{a});
}
array inv(const array& a, StreamOrDevice s /* = {} */) {
array inv_impl(const array& a, bool tri, bool upper, StreamOrDevice s) {
if (a.dtype() != float32) {
std::ostringstream msg;
msg << "[linalg::inv] Arrays must type float32. Received array "
@@ -258,7 +258,21 @@ array inv(const array& a, StreamOrDevice s /* = {} */) {
}
return array(
a.shape(), a.dtype(), std::make_shared<Inverse>(to_stream(s)), {a});
a.shape(),
a.dtype(),
std::make_shared<Inverse>(to_stream(s), tri, upper),
{a});
}
array inv(const array& a, StreamOrDevice s /* = {} */) {
return inv_impl(a, /*tri=*/false, /*upper=*/true, s);
}
array tri_inv(
const array& a,
bool upper /* = true */,
StreamOrDevice s /* = {} */) {
return inv_impl(a, /*tri=*/true, upper, s);
}
array cholesky(
@@ -292,4 +306,37 @@ array cholesky(
{a});
}
array cholesky_inv(
const array& L,
bool upper /* = false */,
StreamOrDevice s /* = {} */) {
if (L.dtype() != float32) {
std::ostringstream msg;
msg << "[linalg::cholesky] Arrays must type float32. Received array "
<< "with type " << L.dtype() << ".";
throw std::invalid_argument(msg.str());
}
if (L.ndim() < 2) {
std::ostringstream msg;
msg << "[linalg::cholesky] Arrays must have >= 2 dimensions. Received array "
"with "
<< L.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
if (L.shape(-1) != L.shape(-2)) {
throw std::invalid_argument(
"[linalg::cholesky] Cholesky inverse is only defined for square "
"matrices.");
}
array L_inv = tri_inv(L, upper, s);
if (upper) {
return matmul(L_inv, swapaxes(L_inv, -1, -2, s), s);
} else {
return matmul(swapaxes(L_inv, -1, -2, s), L_inv, s);
}
}
} // namespace mlx::core::linalg