[Feature request] Improved support for complex64 arrays

[cjhaster]

In my project I need to use complex-valued data. While many operations in mlx.core support complex arrays, there are a few basic examples missing. For example, both:

import mlx.core as mx

test = mx.add(mx.random.uniform(shape=[100]), mx.multiply(1j, mx.random.uniform(shape=[100])))
mx.sum(test)

and

test[0]

causes the kernel to die, whereas if I do these operations on float32 arrays everything works as it should. Another operation I've noticed is not defined for complex64 input arrays is mx.divide if the denominator is complex (it works fine for a complex numerator).

I'm able to get a workaround by casting the mlx arrays to numpy arrays, but this recasting is now the computationally limiting operation in my code 🙂

I'm using MacOS 14.2.1 and mlx 0.7.0 in a Python 3.10 environment.

[cjhaster]

It would also be convenient (and save me another numpy recasting!) if it was possible to access the components of a complex array though test.real and test.imag.

[awni]

Indeed complex64 is not up to standard yet. I think the reduction is not supported for complex! We should definitely fix that and the real/imag ops.

[cjhaster]

I'll fully admit to my C++ knowledge not having been very much used in the last few years, so I don't know how much "instant coding" I'll be able to do to silve this issue for the next release. But since I'm sufficiently interested in these features being added to MLX I'm happy to help out in whatever way I can!