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

Add matrix inversion primitive (#822)

Browse Source
This commit is contained in:
nicolov
2024-03-15 14:34:36 +01:00
committed by GitHub
parent 19ec023256
commit eaba55c9bf
13 changed files with 204 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
mlx/linalg.cpp
View File

@@ -238,4 +238,27 @@ std::vector<array> svd(const array& a, StreamOrDevice s /* = {} */) {
{a});
}
array inv(const array& a, StreamOrDevice s /* = {} */) {
if (a.dtype() != float32) {
std::ostringstream msg;
msg << "[linalg::inv] 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::inv] 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::inv] Inverses are only defined for square matrices.");
}
return array(
a.shape(), a.dtype(), std::make_unique<Inverse>(to_stream(s)), {a});
}
} // namespace mlx::core::linalg