Add determinant and sign-log-determinant functions to mlx.core.linalg#3416
Merged
Conversation
4 tasks
zcbenz
approved these changes
Apr 27, 2026
zcbenz
left a comment
Collaborator
There was a problem hiding this comment.
Looks good to me, thanks!
Can you rebase the branch to fix the flaky test?
Collaborator
|
Also ping @abeleinin who wrote #1451 for a review if possible. |
634025d to
09bb9a8
Compare
abeleinin
reviewed
Apr 28, 2026
LAPACK's getrf returns info > 0 for singular matrices, which is informational (factorization completed, U has zeros on diagonal), not an error. Previously LUF threw on info != 0, blocking det/slogdet from handling singular matrices with n >= 4. Now only throw on info < 0 (illegal argument). Downstream consumers like slogdet already handle singular U correctly (returning sign=0, logabsdet=-inf). Co-Authored-By: Lucas Fernandes Martins <Lucas-Fernandes-Martins@users.noreply.github.com>
ac208e7 to
c7c82cf
Compare
ekryski
added a commit
to ekryski/mlx
that referenced
this pull request
May 6, 2026
…dLocalStream, linalg.det) Highlights: - ThreadLocalStream in C++ (ml-explore#3405) — joins StreamOrDevice variant - Segmented mm NAX kernel (ml-explore#3419) — new steel/gemm/kernels/steel_gemm_segmented{,_nax} - mlx.core.linalg additions: determinant, sign-log-determinant (ml-explore#3416) - NAX split-K tuning + addmm fix (ml-explore#3422) - gather_mm CUDA + qmm CUDA fixes (ml-explore#3414, ml-explore#3417, ml-explore#3445) - Scheduler::enqueue thread safety (ml-explore#3423) - Random kernel bytes_per_key truncation fix (ml-explore#3432) - Many smaller bug + CI fixes; nanobind 2.12.0 pin Conflict resolution (5 files; all auto-merged cleanly, no manual edits): - mlx/utils.h — `ThreadLocalStream` joins StreamOrDevice; our metal_retain_bound_buffers() helper preserved. - mlx/backend/metal/device.cpp — upstream's "no Metal device available" error message coexists with our signpost/MLX_METAL_PROFILE perf diagnostic + hazard-mode retain logic. - mlx/backend/metal/CMakeLists.txt — upstream's steel_gemm_segmented{,_nax} joins our kernel registrations. - mlx/backend/metal/kernels/CMakeLists.txt — same: upstream segmented + our fused_gate_activation kernel. - tests/scheduler_tests.cpp — upstream's "test thread local stream" joins our scheduler-test alignment for lazy command-encoder init. No disputed conflicts; the auto-merge produced a coherent result with both sides' work intact. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes: Add
linalg.detandlinalg.slogdetRelated: #3335
Add determinant (
det) and sign-log-determinant (slogdet) functions tomlx.core.linalg.mx.linalg.det(a)- returns the determinant of a square matrixmx.linalg.slogdet(a)- returns(sign, logabsdet)for numerical stabilityDesign Choices
Composite on
lu_factor(), no new primitiveBoth functions compose on the existing
lu_factor()primitive rather than introducing a new one. This mirrors how JAX implementsslogdetand is consistent with howsolve()already builds onlu()in this codebase.slogdetas the core,detderived from itslogdetis the numerically stable. For the general case (N > 3),detsimply callsslogdetand reconstructs viasign * exp(logabsdet). This avoids duplicating the LU-based logic. An internalslogdet_implfunction handles the shared logic so the publicdetandslogdetentry points each validate and promote types exactly once.Small-matrix fast paths (1x1, 2x2, 3x3)
A shared
det_raw_small()helper computes determinants directly using closed-form formulas (element extraction via slice/squeeze, then the standard cofactor expansion for 3x3). This avoids LU overhead for the small matrices common in batched ML workloads (e.g., 3x3 covariance matrices in normalizing flows).detcallsdet_raw_smalldirectly for N <= 3, skipping the log/exp round-trip entirely.slogdetcalls it too, then splits the raw result intosign(raw)andlog(abs(raw)).Integer auto-promotion
det/slogdetpromote integer inputs to float32 viaat_least_float(). This matches NumPy's behavior (np.linalg.detaccepts integer arrays).Singular matrix handling
For the LU path, singularity is detected by checking for exact zeros on the U diagonal (same approach as NumPy). Singular matrices return
(sign=0, logabsdet=-inf)fromslogdet, and0.0fromdet.Thanks @Lucas-Fernandes-Martins for pointing out the singular matrix issue in LU factorization.
Testing
np.linalg.det/np.linalg.slogdetfor 1x1 through 5x5, random matrices, identity(3, 4, 4)and(2, 3, 3, 3)match NumPy0.1 * I_100(det underflows in float32, slogdet stays finite)sign=0,logabsdet=-infsign * exp(logabsdet)matchesdetfor non-singular inputsValueErrorChecklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes