MLX compatibility: linalg functions

[prady0t]

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

[ev-br]

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]

Yes sure. Let's open a PR there. Thanks.

[ev-br]

Superb! Feel free to ping me on that PR if that's helpful.

[prady0t]

Superb! Feel free to ping me on that PR if that's helpful.

Sorry, I missed your comment. Pinging you now.