Will rearrange be supported in einsum in the future version?

[wzm2256]

I mostly use this lib for tensor computing, and I find it really useful. However, I find that I always need to write several lines of code to implement a simple idea. One example is tensor the product (Kronecker product).

import torch
import einops
a=torch.randn(32,32)
b = einops.einsum(a, a, 'b c, d e -> (b d) (c e)' )
# NotImplementedError: Shape rearrangement is not yet supported in einsum.

# I have to write something like this
b = einops.rearrange(einops.einsum(a, a, 'b c, d e -> b d c e' ), 'b d c e -> (b d) (c e)')

wher I compute b as the kronecker product of a: $b= a\otimes a$, and I feel this is unnecessarily complicated. I think it would be wonderful if rearrange is supported in einsum.

[austinmw]

Would be awesome to support this!

attn_output = einsum(attn_weights, v, "b h s t, b h t d -> b s (h d)")