MLX compatibility: operators and elementwise functions

[ev-br]

Towards #450.

Running

$ ARRAY_API_TESTS_SKIP_DTYPES=float64,complex128 ARRAY_API_TESTS_MODULE=mlx.core pytest array_api_tests/test_operators_and_elementwise_functions.py -v
...
==================================================== short test summary info =====================================================
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_acos - ValueError: a=(nan+nanj) and b=(1.570796326562066-22.873856958478193j), but input must be non-NaN
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_acosh - RuntimeError: [metal::Device] Unable to load kernel v_ArcCoshcomplex64complex64
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_asin - AssertionError: out=infj, but should be roughly asin(x)=(0.7856422146397344+8.664583897600078j) [asin()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_asinh - RuntimeError: [metal::Device] Unable to load kernel v_ArcSinhcomplex64complex64
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_atan - ExceptionGroup: Hypothesis found 2 distinct failures. (2 sub-exceptions)
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_atanh - RuntimeError: [metal::Device] Unable to load kernel v_ArcTanhcomplex64complex64
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[bitwise_left_shift(x1, x2)] - AssertionError: out=1, but should be (x1 << x2)=0 [bitwise_left_shift()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__lshift__(x1, x2)] - AssertionError: out=1, but should be (x1 << x2)=0 [__lshift__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_left_shift[__ilshift__(x1, x2)] - AssertionError: x1=1, but should be (x1 << x2)=0 [__ilshift__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[bitwise_right_shift(x1, x2)] - AssertionError: out=1, but should be (x1 >> x2)=0 [bitwise_right_shift()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__rshift__(x1, x2)] - AssertionError: out=1, but should be (x1 >> x2)=0 [__rshift__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_bitwise_right_shift[__irshift__(x1, x2)] - AssertionError: x1=1, but should be (x1 >> x2)=0 [__irshift__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_clip - TypeError: clip(): incompatible function arguments. The following argument types are supported:
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_copysign - AssertionError: copysign is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_divide[divide(x1, x2)] - ExceptionGroup: Hypothesis found 2 distinct failures. (2 sub-exceptions)
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_expm1 - ValueError: [expm1] Not supported for complex64.
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[floor_divide(x1, x2)] - AssertionError: out=0, but should be (x1 // x2)=-1 [floor_divide()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__floordiv__(x1, x2)] - AssertionError: out=0, but should be (x1 // x2)=-1 [__floordiv__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__ifloordiv__(x1, x2)] - AssertionError: x1=0, but should be (x1 // x2)=-1 [__ifloordiv__()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_hypot - AssertionError: hypot is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_log - AssertionError: out=(inf+0.8922129273414612j), but should be roughly log(x)=(44.361419543366964+0.892212881624599j) [log()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_log1p - AssertionError: out=(inf+0.7905846238136292j), but should be roughly log1p(x)=(44.361419543239215+0.7905845671575851j) [log1p()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_log2 - AssertionError: out=(inf+1.2417556047439575j), but should be roughly log2(x)=(63.99999997183151+1.241755551201958j) [log2()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_log10 - AssertionError: out=(inf+0.47804269194602966j), but should be roughly log10(x)=(19.265919715434027+0.47804270750558625j) [log...
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_nextafter - AssertionError: nextafter is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_positive[__pos__] - AttributeError: 'mlx.core.array' object has no attribute '__pos__'. Did you mean: '__pow__'?
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_signbit - AssertionError: signbit is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt - AssertionError: out=(inf+infj), but should be roughly sqrt(x)=(3956798150.479092+1670476642.4809253j) [sqrt()]
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[atan2] - TypeError: arctan2(): incompatible function arguments. The following argument types are supported:
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[copysign] - AssertionError: copysign is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[hypot] - AssertionError: hypot is not defined in mlx.core
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[nextafter] - AssertionError: nextafter is not defined in mlx.core
===================================== 32 failed, 123 passed, 5 skipped in 126.21s (0:02:06) ======================================

Failures sort into several categories:

• Functions missing in MLX: nextafter, hypot, copysign, signbit, __pos__:
  --Could be submitted to MLX as pure enhancements.
• Functions exist in MLX but miss complex kernels: acosh, asinh, atanh, exmp1:
  --We could report them upstream.
• Early overflow in the complex plane (likely, internal computations done in float32): acos, asin, log, log2, log10, log1p, sqrt:
  --Cross-ref BUG: Platform dependent test failures of tanh, atanh and sqrt function. numpy/numpy#32234 (comment), hopefully will get fixed at the source.
• Functions only arrays and reject python scalars: mx.atan2( mx.asarray(1.0), 2.0) :
  --Looks to be a small omission (other binary functions accept scalars?), can submit a PR upstream.
• Modular arithmetics / rounding modes for integers: __floordiv__(array(-1, dtype=int8), array(2, dtype=int8)), also shifts, bitwise_right_shift(array(1, dtype=int32), array(32, dtype=int16)):
  --These need investigation.
• clip requires min, max arguments, while in the spec they are optional:
  --Also a simple omission, can submit a PR upstream.

[ev-br]

fyi @aaishwarymishra @prady0t

[prady0t]

For me, building MLX with `MLX_BUILD_METAL=OFF ', the complex kernel problem disappears; instead we have:

FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_acosh - AssertionError: out=(inf+0j), but should be roughly acosh(x)=(9.011157473000706+3.141348572552466j) [acosh()]

FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_asinh - AssertionError: out=(-16.63553237915039+1.5707963705062866j), but should be roughly asinh(x)=(-8.605570402591917+0.00036616619856045374j) [asinh()]

atanh: pass!

exmp1: not sure what it means, not in your failure log either.

[ev-br]

That's a nice data point indeed.

As far as I understand, MLX carries at least three separate implementations of the core maths: CPU, CUDA and Metal. It's not very surprising that they all have their own shortcomings, and anecdotally, we've seen with pytorch that the Metal one has more rough edges than others. Since the main push, also the default mode for MLX is Metal, I'd focus on things which are general for all backends (if this is the term here), at least for the time being.