Support for einops

[gboduljak]

Hi :)

I am a big fan of einops and I believe that supporting it within the framework would be highly valuable, especially for quick prototyping and situations where performance isn't the primary focus. Are there any plans to implement a feature akin to torch.einsum within the framework? If there's a plan for something like mlx.einsum and contributions are welcome, I'd be eager to contribute.

[gboduljak]

It seems that Apple MLX support was already proposed within einops.

[sck-at-ucy]

I second that. einops saved me quite a few times where the logic was becoming a bit intractable.

[awni]

We'll definitely keep this on the feature roadmap. I think if we had an einsum like NumPY integration with einops is doable. (Contributions welcome!)

[dc-dc-dc]

Will also need a tile op to support einops

[dc-dc-dc]

After #278 and #240 get merged, that should be all thats needed to create a backend for einops

[arogozhnikov]

Einops maintainer here.

Simplest way to get einops functions for framework is to implement array api.
einops.array_api works with all array-api compatible tensors (there is also support for scikit-learn and eindex).

It also simplifies prioritization of ops a lot for framework devs, as what's in the standard is a bare minimum.

[awni]

We can close this as the latest MLX works with einops.array_api back-end.

[ifsheldon]

Hi! I'm running this code, trying to use einops.einsum with mx.array, but it doesn't work. So, einsum is still not supported?

a = mx.random.normal([2, 3])
b = mx.random.normal([4, 5])

kron1 = mx.kron(a, b)
kron2 = einsum(a, b, "a b, c d -> a c b d")
kron2 = rearrange(kron2, "a c b d -> (a c) (b d)")

RuntimeError                              Traceback (most recent call last)
Cell In[17], line 5
      2 b = mx.random.normal([4, 5])
      4 kron1 = mx.kron(a, b)
----> 5 kron2 = einsum(a, b, "a b, c d -> a c b d")
      6 kron2 = rearrange(kron2, "a c b d -> (a c) (b d)")
      8 assert mx.allclose(kron1, kron2)

File ~/offline_code/personal/tensor_network/.venv/lib/python3.12/site-packages/einops/einops.py:916, in einsum(*tensors_and_pattern)
    914 tensors = tensors_and_pattern[:-1]
    915 pattern = _compactify_pattern_for_einsum(pattern)
--> 916 return get_backend(tensors[0]).einsum(pattern, *tensors)

File ~/offline_code/personal/tensor_network/.venv/lib/python3.12/site-packages/einops/_backends.py:59, in get_backend(tensor)
     56                 _type2backend[_type] = backend
     57                 return backend
---> 59 raise RuntimeError("Tensor type unknown to einops {}".format(type(tensor)))

RuntimeError: Tensor type unknown to einops <class 'mlx.core.array'>

[awni]

It looks like the einops.array_api package doesn't have einsm, so you'd need to use MLXs einsum and then the rearrange from einops.array_api that like so:

import mlx.core as mx
from einops.array_api import rearrange
a = mx.random.normal([2, 3])
b = mx.random.normal([4, 5])

kron1 = mx.kron(a, b)
kron2 = mx.einsum("a b, c d -> a c b d", a, b)
kron2 = rearrange(kron2, "a c b d -> (a c) (b d)")

[awni]

I'd recommend filing an issue on einops github to have better builtin support for MLX and/or support MLX arrays in einsum.

[ifsheldon]

OK, I filed a PR to that repo a few hours ago. Thanks!