Repository navigation
MLX compatibility: linalg functions #473
Description
Scroll ->
| Array API | MLX Analog | Status | Notes | Test Node | Result |
|---|---|---|---|---|---|
| linalg.cholesky(x, /, *, upper=False) | linalg.cholesky(a, /, *, upper: bool = False, stream=None) | incompatible | MLX rejects complex64 inputs and only supports float32 and float64, while the test generates supported complex inputs. |
test_cholesky | Failed |
| linalg.det(x, /) | linalg.det(a, /, *, stream=None) | incompatible | MLX does not support complex inputs; complex64 raises ValueError. |
test_det | Failed |
| linalg.diagonal(x, /, *, offset=0) | No MLX analog | missing | mlx.core.linalg has no diagonal attribute. |
test_diagonal | Failed |
| linalg.eigh(x, /) | linalg.eigh(a, /, *, stream=None) | incompatible | Multiple issues: MLX fails while generating some batched symmetric matrices due to broadcasting, rejects empty/non-square matrix cases expected by the tests, and returns a plain tuple instead of the required namedtuple with eigenvalues and eigenvectors. |
test_eigh | Failed |
| linalg.eigvalsh(x, /) | linalg.eigvalsh(a, /, *, stream=None) | incompatible | MLX fails for some batched symmetric matrices due to broadcasting and rejects empty/non-square matrix cases generated by the tests. | test_eigvalsh | Failed |
| linalg.eig(x, /) | linalg.eig(a, /, *, stream=None) | incompatible | MLX returns a plain tuple instead of the required namedtuple with eigenvalues and eigenvectors. |
test_eig | Failed |
| linalg.inv(x, /) | linalg.inv(a, /, *, stream=None) | incompatible | MLX rejects complex64 inputs and only supports float32 and float64. |
test_inv | Failed |
| linalg.matmul(x1, x2, /) | No mlx.core.linalg.matmul analog |
missing | mlx.core.linalg has no matmul attribute. |
test_linalg_matmul | Failed |
| matmul(x1, x2, /) | matmul(a, b, /, *, stream=None) | incompatible | MLX rejects integer inputs such as uint16, requiring inexact/floating-point result types. |
test_matmul | Failed |
| linalg.matrix_norm(x, /, *, keepdims=False, ord="fro") | No MLX analog | missing | mlx.core.linalg has no matrix_norm attribute. |
test_matrix_norm | Failed |
| linalg.matrix_power(x, n, /) | No MLX analog | missing | mlx.core.linalg has no matrix_power attribute. |
test_matrix_power | Failed |
| linalg.matrix_rank(x, /, *, rtol=None) | No MLX analog | missing | mlx.core.linalg has no matrix_rank attribute. |
test_matrix_rank | Failed |
| linalg.matrix_transpose(x, /) | No mlx.core.linalg.matrix_transpose analog |
missing | mlx.core.linalg has no matrix_transpose attribute. |
test_linalg_matrix_transpose | Failed |
| matrix_transpose(x, /) | matrix_transpose(a, /, *, stream=None) | incompatible | MLX produces an incorrect output shape for empty arrays with rank greater than two; (0, 1, 1) becomes (1, 1, 0) instead of swapping only the last two axes. |
test_matrix_transpose | Failed |
| linalg.outer(x1, x2, /) | No MLX analog | missing | mlx.core.linalg has no outer attribute. |
test_outer | Failed |
| linalg.pinv(x, /, *, rtol=None) | linalg.pinv(a, /, *, stream=None) | incompatible | MLX lacks the rtol keyword and also rejects complex64 inputs, supporting only float32 and float64. |
test_pinv | Failed |
| linalg.qr(x, /, *, mode="reduced") | linalg.qr(a, /, *, stream=None) | incompatible | MLX lacks the mode keyword, returns a plain tuple instead of a namedtuple with Q and R, and rejects complex64 inputs. |
test_qr | Failed |
| linalg.slogdet(x, /) | linalg.slogdet(a, /, *, stream=None) | incompatible | MLX returns a plain tuple instead of the required namedtuple with sign and logabsdet, and does not support complex inputs. |
test_slogdet | Failed |
| linalg.solve(x1, x2, /) | linalg.solve(a, b, /, *, stream=None) | incompatible | MLX fails on an empty-array case with a take_along_axis dimension mismatch. |
test_solve | Failed |
| linalg.svd(x, /, *, full_matrices=True) | linalg.svd(a, compute_uv: bool = True, *, stream=None) | incompatible | MLX exposes compute_uv instead of the Array API's full_matrices keyword and returns a plain tuple instead of a namedtuple with U, S, and Vh. |
test_svd | Failed |
| linalg.svdvals(x, /) | No MLX analog | missing | mlx.core.linalg has no svdvals attribute. |
test_svdvals | Failed |
| linalg.tensordot(x1, x2, /, *, axes=2) | No mlx.core.linalg.tensordot analog |
missing | mlx.core.linalg has no tensordot attribute. |
test_linalg_tensordot | Failed |
| tensordot(x1, x2, /, *, axes=2) | tensordot(a, b, /, *, axes=2, stream=None) | incompatible | MLX fails for integer inputs because its implementation reaches matmul, which only supports inexact/floating-point types. |
test_tensordot | Failed |
| linalg.trace(x, /, *, offset=0, dtype=None) | No MLX analog | missing | mlx.core.linalg has no trace attribute. |
test_trace | Failed |
| linalg.vecdot(x1, x2, /, *, axis=-1) | No mlx.core.linalg.vecdot analog |
missing | mlx.core.linalg has no vecdot attribute. |
test_linalg_vecdot | Failed |
| vecdot(x1, x2, /, *, axis=-1) | vecdot(a, b, /, *, axis=-1, stream=None) | incompatible | MLX produces an incorrect promoted dtype for mixed integer inputs: uint8 and uint16 produce uint32, while the expected dtype is uint16. |
test_vecdot | Failed |
| linalg.vecdot(x1, x2, /, *, axis=-1) | No mlx.core.linalg.vecdot analog |
missing | mlx.core.linalg.vecdot is missing, so the required conjugating vecdot behavior cannot be tested. |
test_vecdot_conj | Failed |
| linalg.vector_norm(x, /, *, axis=None, keepdims=False, ord=2) | No MLX analog | missing | mlx.core.linalg has no vector_norm attribute. |
test_vector_norm | Failed |
Activity
[-]MLX Compatablilty: linalg functions[/-][+]MLX compatiblilty: linalg functions[/+][-]MLX compatiblilty: linalg functions[/-][+]MLX compatibility: linalg functions[/+]ev-br commentedon Aug 30, 2026
Thanks @prady0t . The picture is pretty clear: there's a mix of missing functions (relatively easy to add in principle), missing edge cases, and missing low-level kernels (e.g. complex). The latter can be very tricky to implement.
Taking a step back, I think it might be about time to work on ml-explore/mlx#3514 so that we share a general overview of compatibility with MLX developers, and can prioritize defects with them.
Is this somethign you'd be interested to work on?
I think it would be best to have the test suite run on a metal CI, skipping float64 and complex128, and using
--disable-data-dependent-shapes, and see the full list of failures. For skips, special case failures can be skipped in bulk and are somewhat low prio; other failures we'll be able to work through with upstream devs.prady0t commentedon Aug 30, 2026
Yes sure. Let's open a PR there. Thanks.
ev-br commentedon Aug 30, 2026
Superb! Feel free to ping me on that PR if that's helpful.
prady0t commentedon Aug 30, 2026
Sorry, I missed your comment. Pinging you now.