Add support for computing matrix determinants

[jessegrabowski]

The mlx.core.linalg sub-module doesn't have any functions related to computation of determinants. These pop up all over the place, as part of the lop of a multivariate normal in PPL land, but also when compute change-of-variable correction terms in normalizing flow layers. Would be nice to have at least mx.linalg.det. A numerically stable mx.linalg.slogdet is also important.

[abhilashreddys]

@zcbenz I'd like to take this on. I plan to implement both slogdet and det using the existing lu_factor decomposition.

[Lucas-Fernandes-Martins]

Hi @abhilashreddys I've been working on the same thing for the past few days. If you want to take a look at my PR maybe we can collaborate on this one. For now, I have only implemented linalg.det so there's still slogdet to be done. Also if you have any suggestions on my implementation I'd love to hear your thoughs :)

[abhilashreddys]

Hi @zcbenz @Lucas-Fernandes-Martins I was also working on it for the past few days. I have implemented mx.linalg.det and mx.linalg.slogdet as composites on the existing lu_factor(). slogdet is the core implementation, and det is derived from it via sign * exp(logabsdet), so it benefits from the same numerical stability. Also includes small-matrix fast paths (1x1/2x2/3x3) for batched workloads and full test coverage against NumPy.

PR: #3416

Let me know your thoughts!

[Lucas-Fernandes-Martins]

@abhilashreddys nice, we can bounce off each other then :)

I'm testing your implementation and it doesn't seem to work for singular matrices which have 4 or more dimensions. I believe that's because there's some error handling on [mlx/backend/cpu/luf.cpp](https://github.com/ml-explore/mlx/pull/3415/changes#diff-e4ba8f7bef27f49a7d15aef13ba877301098fdebffbf015af2998e3641dfb1ab) that throws an error if the matrix is singular. That's why I introduced the allow_singular boolean flag on lu_helper.

Let me know your opinion on this

[abhilashreddys]

@Lucas-Fernandes-Martins, thanks for pointing out this issue. I have implemented a fix (added you as a co-author in the commit). I changed lines 70-76 of luf.cpp to only throw on info < 0 (illegal argument). Let info > 0 (singular) pass through silently. This matches NumPy/SciPy's getrf behavior.

In your PR, you added an optional allow_singular parameter to LUF, with a default of false for backward compatibility. But every downstream consumer (solve, inv) already handles singularity. Do you have any thoughts on why adding a flag is better?

[Lucas-Fernandes-Martins]

@abhilashreddys thank you :)

So I considered both approaches when implementing. I agree that changing the luf.cpp throw condition is the cleaner solution (specially since it matches what Numpy does).

I added the flag not to change the behavior of any other function in the API. Not throwing for 'info > 0' would make the LU factorization not throw an error even for matrices that do not admit a proper LU factorization.

I agree with you that Numpy also not throws for info > 0, but it allows Scipy to add a Reconstruction error: True warning since it provides the _raise_linalgerror_singular callback.

Currently, MLX doesn't seem to have a similar callback mechanism. Therefore, I think it is non-ideal to not throw on info > 0 and leave the user with a degenerate factorization and no error message.

Adding the flag was a non-breaking way to let linalg.det call lu_helper without throwing an error, but keeping the errors thrown when linalg.lu_factor is called. My idea was to add linalg.det without changing the behavior of any other method. I agree the ideal solution would be to implement this in the Python wrapper, similar to how Numpy does it, but this would touch a lot of the code.

@zcbenz if you have any hints or advice on how to do this better :)

[abhilashreddys]

@Lucas-Fernandes-Martins Good point about the silent degenerate factorization. How about we keep the LUF primitive clean (no throw for info > 0, matching NumPy's getrf) but add a singularity check at the lu_factor Python wrapper level? This way det/slogdet work correctly, and lu_factor can still warn users about singular matrices. The error boundary moves from the LAPACK primitive to the API layer, which is where NumPy/SciPy handle it too.

[jessegrabowski]

fwiw jax/xla does not warn or raise on linalg errors, which is a highly desirable behavior when running large models, potentially random inputs (e.g. MCMC, minibatch training). Most of these use-cases are robust to a single failed attempt. Having the entire program halt due to numerical instability is a big ick. Same with an avalanche of warning spam. Better to trust users give good inputs and write good programs with their own error handling.

[abhilashreddys]

Thanks, @jessegrabowski , for the thoughtful feedback on the flag approach.

The one-line change in luf.cpp with no new parameters, no flags, and no behavior change for any consumer except that they now receive valid output instead of an exception. Should resolve all the issues that are pointed out. The PR I raised is complete as-is.

[zcbenz]

I wonder if @abeleinin who wrote #1451 could help review the PRs, my knowledge with linear algebra is quite rusty.

[Lucas-Fernandes-Martins]

@jessegrabowski @abhilashreddys if we're changing the behaviour of luf.cpp (which will impact what throws or not in lu_solve for instance) maybe we should break this into two PRs: one where the behaviour of luf.cpp is changed, and then the det one. Or maybe just make it clear on the PR this will change the current behaviour of lu_factor regarding singular matrices.

[Lucas-Fernandes-Martins]

@angeloskath do you have any thoughts on what would be the correct philosophy about error handling for methods such as lu_factor? Would it make sense to change the behavior of luf.cpp so we only treat the singularity of the matrix upstream, instead of throwing an error straight away?