MLX - potential adoption?

[lucascolley]

https://github.com/ml-explore/mlx

How close does MLX look to the standard? Has there been any discussion about this so far?

[rgommers]

Very interesting, thanks for sharing @lucascolley. It looks reasonably close, since it's modeled on NumPy. There hasn't been any discussion about it yet. Based on the issues, which were all opened in the last 24 hours, this was just publicly launched.

The thing that stands out to me is unified memory, allowing execution on GPU without having to move the data. That's a very interesting capability, which may also have impact on the UX of the library.

[asmeurer]

It looks like the biggest issue is that it doesn't have float64. That makes it very hard to run the test suite on it, at least. I was able to run the has_names test (see data-apis/array-api-tests#223) and got

========================================================================= short test summary info =========================================================================
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-cholesky] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-cross] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-det] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-diagonal] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-eigh] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-eigvalsh] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-inv] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-matmul] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-matrix_norm] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-matrix_power] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-matrix_rank] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-matrix_transpose] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-outer] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-pinv] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-qr] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-slogdet] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-solve] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-svd] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-svdvals] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-tensordot] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-trace] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-vecdot] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[linalg-vector_norm] - AttributeError: module 'mlx.core' has no attribute 'linalg'
FAILED array_api_tests/test_has_names.py::test_has_names[statistical-std] - AssertionError: mlx.core is missing the statistical function std()
FAILED array_api_tests/test_has_names.py::test_has_names[set-unique_all] - AssertionError: mlx.core is missing the set function unique_all()
FAILED array_api_tests/test_has_names.py::test_has_names[set-unique_counts] - AssertionError: mlx.core is missing the set function unique_counts()
FAILED array_api_tests/test_has_names.py::test_has_names[set-unique_inverse] - AssertionError: mlx.core is missing the set function unique_inverse()
FAILED array_api_tests/test_has_names.py::test_has_names[set-unique_values] - AssertionError: mlx.core is missing the set function unique_values()
FAILED array_api_tests/test_has_names.py::test_has_names[searching-nonzero] - AssertionError: mlx.core is missing the searching function nonzero()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-asarray] - AssertionError: mlx.core is missing the creation function asarray()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-empty] - AssertionError: mlx.core is missing the creation function empty()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-empty_like] - AssertionError: mlx.core is missing the creation function empty_like()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-eye] - AssertionError: mlx.core is missing the creation function eye()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-from_dlpack] - AssertionError: mlx.core is missing the creation function from_dlpack()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-full_like] - AssertionError: mlx.core is missing the creation function full_like()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-linspace] - AssertionError: mlx.core is missing the creation function linspace()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-meshgrid] - AssertionError: mlx.core is missing the creation function meshgrid()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-tril] - AssertionError: mlx.core is missing the creation function tril()
FAILED array_api_tests/test_has_names.py::test_has_names[creation-triu] - AssertionError: mlx.core is missing the creation function triu()
FAILED array_api_tests/test_has_names.py::test_has_names[manipulation-concat] - AssertionError: mlx.core is missing the manipulation function concat()
FAILED array_api_tests/test_has_names.py::test_has_names[manipulation-flip] - AssertionError: mlx.core is missing the manipulation function flip()
FAILED array_api_tests/test_has_names.py::test_has_names[manipulation-permute_dims] - AssertionError: mlx.core is missing the manipulation function permute_dims()
FAILED array_api_tests/test_has_names.py::test_has_names[manipulation-roll] - AssertionError: mlx.core is missing the manipulation function roll()
FAILED array_api_tests/test_has_names.py::test_has_names[manipulation-stack] - AssertionError: mlx.core is missing the manipulation function stack()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-astype] - AssertionError: mlx.core is missing the data_type function astype()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-broadcast_arrays] - AssertionError: mlx.core is missing the data_type function broadcast_arrays()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-can_cast] - AssertionError: mlx.core is missing the data_type function can_cast()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-finfo] - AssertionError: mlx.core is missing the data_type function finfo()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-iinfo] - AssertionError: mlx.core is missing the data_type function iinfo()
FAILED array_api_tests/test_has_names.py::test_has_names[data_type-result_type] - AssertionError: mlx.core is missing the data_type function result_type()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-acos] - AssertionError: mlx.core is missing the elementwise function acos()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-acosh] - AssertionError: mlx.core is missing the elementwise function acosh()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-asin] - AssertionError: mlx.core is missing the elementwise function asin()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-asinh] - AssertionError: mlx.core is missing the elementwise function asinh()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-atan] - AssertionError: mlx.core is missing the elementwise function atan()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-atan2] - AssertionError: mlx.core is missing the elementwise function atan2()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-atanh] - AssertionError: mlx.core is missing the elementwise function atanh()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_and] - AssertionError: mlx.core is missing the elementwise function bitwise_and()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_left_shift] - AssertionError: mlx.core is missing the elementwise function bitwise_left_shift()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_invert] - AssertionError: mlx.core is missing the elementwise function bitwise_invert()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_or] - AssertionError: mlx.core is missing the elementwise function bitwise_or()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_right_shift] - AssertionError: mlx.core is missing the elementwise function bitwise_right_shift()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-bitwise_xor] - AssertionError: mlx.core is missing the elementwise function bitwise_xor()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-ceil] - AssertionError: mlx.core is missing the elementwise function ceil()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-expm1] - AssertionError: mlx.core is missing the elementwise function expm1()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-floor] - AssertionError: mlx.core is missing the elementwise function floor()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-floor_divide] - AssertionError: mlx.core is missing the elementwise function floor_divide()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-isfinite] - AssertionError: mlx.core is missing the elementwise function isfinite()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-isinf] - AssertionError: mlx.core is missing the elementwise function isinf()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-isnan] - AssertionError: mlx.core is missing the elementwise function isnan()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-logical_and] - AssertionError: mlx.core is missing the elementwise function logical_and()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-logical_or] - AssertionError: mlx.core is missing the elementwise function logical_or()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-logical_xor] - AssertionError: mlx.core is missing the elementwise function logical_xor()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-positive] - AssertionError: mlx.core is missing the elementwise function positive()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-pow] - AssertionError: mlx.core is missing the elementwise function pow()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-remainder] - AssertionError: mlx.core is missing the elementwise function remainder()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-round] - AssertionError: mlx.core is missing the elementwise function round()
FAILED array_api_tests/test_has_names.py::test_has_names[elementwise-trunc] - AssertionError: mlx.core is missing the elementwise function trunc()
FAILED array_api_tests/test_has_names.py::test_has_names[linear_algebra-matrix_transpose] - AssertionError: mlx.core is missing the linear_algebra function matrix_transpose()
FAILED array_api_tests/test_has_names.py::test_has_names[linear_algebra-tensordot] - AssertionError: mlx.core is missing the linear_algebra function tensordot()
FAILED array_api_tests/test_has_names.py::test_has_names[linear_algebra-vecdot] - AssertionError: mlx.core is missing the linear_algebra function vecdot()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__abs__] - AssertionError: The mlx.core array object is missing the method __abs__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__and__] - AssertionError: The mlx.core array object is missing the method __and__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__array_namespace__] - AssertionError: The mlx.core array object is missing the method __array_namespace__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__dlpack__] - AssertionError: The mlx.core array object is missing the method __dlpack__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__dlpack_device__] - AssertionError: The mlx.core array object is missing the method __dlpack_device__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__float__] - AssertionError: The mlx.core array object is missing the method __float__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__floordiv__] - AssertionError: The mlx.core array object is missing the method __floordiv__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__index__] - AssertionError: The mlx.core array object is missing the method __index__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__int__] - AssertionError: The mlx.core array object is missing the method __int__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__invert__] - AssertionError: The mlx.core array object is missing the method __invert__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__lshift__] - AssertionError: The mlx.core array object is missing the method __lshift__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__mod__] - AssertionError: The mlx.core array object is missing the method __mod__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__or__] - AssertionError: The mlx.core array object is missing the method __or__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__pos__] - AssertionError: The mlx.core array object is missing the method __pos__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__rshift__] - AssertionError: The mlx.core array object is missing the method __rshift__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-__xor__] - AssertionError: The mlx.core array object is missing the method __xor__()
FAILED array_api_tests/test_has_names.py::test_has_names[array_method-to_device] - AssertionError: The mlx.core array object is missing the method to_device()
FAILED array_api_tests/test_has_names.py::test_has_names[array_attribute-device] - AssertionError: The mlx.core array object is missing the attribute device
FAILED array_api_tests/test_has_names.py::test_has_names[array_attribute-mT] - AssertionError: The mlx.core array object is missing the attribute mT

So there are also some missing functions and methods, and some places where they use numpy names instead of array API ones.

[leofang]

Yes if they use Metal then they have very limited support for numerical types (ex: no double)

[asmeurer]

I wonder if we should have more concessions for missing dtypes. The standard says they all have to be implemented but in practice quite a few libraries or devices are missing some (at any rate, even if we don't change the spec, the test suite should be lenient about this).

[leofang]

Aren't we (@oleksandr-pavlyk) already pushing for missing dtype support? What's missing that we haven't captured? I am aware we still need one or two PRs to finalize it.

[leofang]

Thank you for bring the new array library to our attention, @lucascolley, though I don't think there's anything actionable on our (Data API Consortium) side. @lucascolley if you can bring array API standard to the MLX folks attention, it'd be great. (Even better is to have an MLX person come and share their feedbacks with us, should they consider adopting the standard.)

I suggest to close this issue and move the remaining discussions (test updates, missing dtypes, ...) to perspective issues.

[asmeurer]

Aren't we (@oleksandr-pavlyk) already pushing for missing dtype support? What's missing that we haven't captured? I am aware we still need one or two PRs to finalize it.

I don't think there's an issue for it yet (there's #672 but that's a little different).

[betatim]

👍 for continuing the discussion around missing dtypes (probably in #640).

PyTorch has the MPS device, which lets you use the GPU builtin to modern MacBooks and there you already have the issue of missing float64.

I think the idea of unified memory is something we should get used to and think about, in terms of what this means for the Array API standard. The result of thinking and coming to a conclusion on this would probably be a few edits to https://data-apis.org/array-api/latest/design_topics/device_support.html

[oleksandr-pavlyk]

Aren't we (@oleksandr-pavlyk) already pushing for missing dtype support? What's missing that we haven't captured? I am aware we still need one or two PRs to finalize it.

Yes, the reality is that some implementations may not provide support for certain data types mandated by the standard, be it for a particular class of devices (MPS, or dpctl.tensor on Iris Xe iGPU not natively supporting xp.float64), or uniformly for all devices (such as PyTorch not supporting unsigned integral types).

The standard should not consider such data type support gaps violations of the spec in my opinion, but must equip end-users with ways to check for required support. In case of dpctl.tensor and oneAPI this might amount to selecting a device with required capabilities, and failing should there be none. For certain workflows reduced datatype support may be sufficient.

[lucascolley]

I suggest to close this issue and move the remaining discussions (test updates, missing dtypes, ...) to perspective issues.

👍 (see ml-explore/mlx#48 for further updates)