zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 01:50:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Enforce triangular matrix form in tri_inv (#1876)

Browse Source 
* fix tri_inv bug

* Revert "fix tri_inv bug"

This reverts commit b74b290201.

* Make sure that tri_inv returns a triangular matrix

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Abe Leininger
2025-02-19 14:42:33 -06:00
committed by GitHub
parent 71de73a668
commit 344a29506e
2 changed files with 23 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
mlx/backend/cpu/inverse.cpp
View File

@@ -81,7 +81,22 @@ void general_inv(array& inv, int N, int i) {
void tri_inv(array& inv, int N, int i, bool upper) {
const char uplo = upper ? 'L' : 'U';
const char diag = 'N';
int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
float* data = inv.data<float>() + N * N * i;
int info = strtri_wrapper(uplo, diag, data, N);
// zero out the other triangle
if (upper) {
for (int i = 0; i < N; i++) {
std::fill(data, data + i, 0.0f);
data += N;
}
} else {
for (int i = 0; i < N; i++) {
std::fill(data + i + 1, data + N, 0.0f);
data += N;
}
}
if (info != 0) {
std::stringstream ss;
ss << "inverse_impl: triangular inversion failed with error code " << info;