[Feature] Ellipsis for broadcasting with einsum

[ashegde]

I noticed that using mlx.core.einsum with ellipsis produces an error. In numpy, I've found that ellipses with np.einsum can be quite useful. Are there plans of including this feature in the future?

To reproduce
MLX code:

import mlx.core as mx

x = mx.random.normal(shape=(3, 4, 5, 6))
y = mx.random.normal(shape=(3, 4, 5, 6))

z = mx.einsum('...ij,...kj->...ik', x, y)

and the resulting error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], [line 7]
      [4] x = mx.random.normal(shape=(3, 4, 5, 6))
      [5] y = mx.random.normal(shape=(3, 4, 5, 6))
----> [7] z = mx.einsum('...ij,...kj->...ik', x, y)

ValueError: [einsum] Subscripts must be letters, but got '.'.

In contrast, the following code works without issue in numpy:

import numpy as np

x = np.random.normal(size=(3, 4, 5, 6))
y = np.random.normal(size=(3, 4, 5, 6))

z = np.einsum('...ij,...kj->...ik', x, y)

Thanks!

[awni]

Indeed we should add support for ellipsis in the subscripts.