Skip to content

Add determinant and sign-log-determinant functions to mlx.core.linalg#3416

Merged
zcbenz merged 7 commits into
ml-explore:mainfrom
abhilashreddys:feat/det
May 5, 2026
Merged

Add determinant and sign-log-determinant functions to mlx.core.linalg#3416
zcbenz merged 7 commits into
ml-explore:mainfrom
abhilashreddys:feat/det

Conversation

@abhilashreddys

@abhilashreddys abhilashreddys commented Apr 15, 2026

Copy link
Copy Markdown
Contributor

Proposed changes: Add linalg.det and linalg.slogdet

Related: #3335

Add determinant (det) and sign-log-determinant (slogdet) functions to mlx.core.linalg.

  • mx.linalg.det(a) - returns the determinant of a square matrix
  • mx.linalg.slogdet(a) - returns (sign, logabsdet) for numerical stability

Design Choices

Composite on lu_factor(), no new primitive

Both functions compose on the existing lu_factor() primitive rather than introducing a new one. This mirrors how JAX implements slogdet and is consistent with how solve() already builds on lu() in this codebase.

slogdet as the core, det derived from it

slogdet is the numerically stable. For the general case (N > 3), det simply calls slogdet and reconstructs via sign * exp(logabsdet). This avoids duplicating the LU-based logic. An internal slogdet_impl function handles the shared logic so the public det and slogdet entry 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).

det calls det_raw_small directly for N <= 3, skipping the log/exp round-trip entirely. slogdet calls it too, then splits the raw result into sign(raw) and log(abs(raw)).

Integer auto-promotion

det/slogdet promote integer inputs to float32 via at_least_float(). This matches NumPy's behavior (np.linalg.det accepts 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) from slogdet, and 0.0 from det.

Thanks @Lucas-Fernandes-Martins for pointing out the singular matrix issue in LU factorization.

Testing

  • Correctness: Compared against np.linalg.det / np.linalg.slogdet for 1x1 through 5x5, random matrices, identity
  • Batched: Shapes (3, 4, 4) and (2, 3, 3, 3) match NumPy
  • Numerical stability: 0.1 * I_100 (det underflows in float32, slogdet stays finite)
  • Singular matrices: sign=0, logabsdet=-inf
  • Consistency: sign * exp(logabsdet) matches det for non-singular inputs
  • Dtypes: float32, float64, integer auto-promotion
  • Error handling: non-square and 1D inputs raise ValueError
  • Full suite: All 18 existing linalg tests continue to pass (186 subtests)

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, thanks!

Can you rebase the branch to fix the flaky test?

Comment thread mlx/linalg.cpp
@zcbenz

zcbenz commented Apr 27, 2026

Copy link
Copy Markdown
Collaborator

Also ping @abeleinin who wrote #1451 for a review if possible.

@abeleinin abeleinin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR looks good! Added some comments about edge cases. @zcbenz Feel free to veto any of my suggestions.

Comment thread mlx/linalg.cpp Outdated
Comment thread python/tests/test_linalg.py
abhilashreddys and others added 6 commits April 28, 2026 11:46
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>
@zcbenz zcbenz merged commit 0938db7 into ml-explore:main May 5, 2026
16 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants