Function lu 

pub fn lu(a: impl AsRef<Array>) -> Result<(Array, Array, Array)>

Compute the LU factorization of the given matrix A.

Note, unlike the default behavior of scipy.linalg.lu, the pivots are indices. To reconstruct the input use L[P, :] @ U for 2 dimensions or mx.take_along_axis(L, P[…, None], axis=-2) @ U for more than 2 dimensions.

To construct the full permutation matrix do:

use mlx_rs::{array, linalg::lu, ops, with_stream, Stream};
use mlx_rs::ops::indexing::{IndexOp, Ellipsis, NewAxis};

with_stream(&Stream::cpu(), || {
    let a = array!([[3.0f32, 1.0, 2.0], [1.0, 8.0, 6.0], [9.0, 2.0, 5.0]]);
    let (p, l, u) = lu(&a).unwrap();
    let p = ops::indexing::put_along_axis(
        mlx_rs::ops::zeros_like(&l).unwrap(),
        p.index((Ellipsis, NewAxis)),
        array!(1.0),
        -1,
    ).unwrap();
    let reconstructed = p.matmul(&l).unwrap().matmul(&u).unwrap();
    assert!(ops::all_close(&reconstructed, &a, None, None, None).unwrap());
});

Params

• a: input array
• stream: stream to execute the operation

Returns

The p, L, and U arrays, such that A = L[P, :] @ U