[Feature suggestion] dimension name alias -> easier tensor contraction

Hi! I'm learning quantum computing and tensor network, so I need to do a lot of tensor contractions. Thank you for developing such a great library. It really makes my life easier.

I'd like to propose a feature called dimension name alias, which should hopefully make writing tensor contraction operations even more easier.

Use case

Here is the problem I faced and below is the illustration.

I have 2 tensors in difference Hilbert spaces and I need to make an outer product of them. If you are familiar with quantum computing or tensor network, these are coefficient tensors for two quantum gates. So, I have 4 qubits. The first gate a acts on first 3 qubits and the second gate b acts on the last 2 qubits. I need to calculate the ensemble of a and b, a coefficient tensor (shape (2, 2, 2, 2, 2, 2, 2, 2)) of a 4-qubit gate.

Here is my demo code

import torch
from einops import einsum

# a is in space H1 x H2 x H3
a = torch.randn(2, 2, 2, # a_l1, a_l2, a_l3
                2, 2, 2, # a_r1, a_r2, a_r3
                )

I = torch.eye(2)
# A is in space H1 x H2 x H3 x H4 by outer product of a and I
A = einsum(a, I, "al1 al2 al3 ar1 ar2 ar3, i4l i4r -> al1 al2 al3 i4l ar1 ar2 ar3 i4r")

# b is in space H3 x H4
b = torch.randn(2, 2, # b_l3, b_l4
                2, 2, # b_r3, b_r4
                )
# B is in space H1 x H2 x H3 x H4 by outer product of two identity gates
B = einsum(I, I, b, "i1l i1r, i2l i2r, bl3 bl4 br3 br4 -> i1l i2l bl3 bl4 i1r i2r br3 br4")

# BA is in space H1 x H2 x H3 x H4 by tensor product of B and A
BA = einsum(B, A, "i1l i2l bl3 bl4 c1 c2 c3 c4, c1 c2 c3 c4 ar1 ar2 ar3 i4r -> i1l i2l bl3 bl4 ar1 ar2 ar3 i4r")
# ba is in space H1 x H2 x H3 x H4 by tensor product of b and a
ba = einsum(b, a, "bl3 bl4 c br4, al1 al2 c ar1 ar2 ar3 -> al1 al2 bl3 bl4 ar1 ar2 ar3 br4")
assert torch.allclose(BA, ba)

Notice that in the lines that calculate BA and ba, I need to keep track of the common dimensions (i.e., the ones to be contracted) myself, like c1, c2 etc, but I just wanted to write down the dimension names first and then specify which dimensions to contract.

Solution

Here comes dimension name alias. The code probably looks like this. I haven't put much thought on the API design.

# einsum can generally accept a list of list of dimensions names that are alias to each other. A list of names is also acceptable.
ba = einsum(b, a, "bl3 bl4 br3 br4, al1 al2 al3 ar1 ar2 ar3 -> al1 al2 bl3 bl4 ar1 ar2 ar3 br4", alias=[
    ["br3", "al3"]
])

I think this can be easily implemented with string manipulation. Just replace that alias names with a unique name. But probably we should make another function instead of making einsum too overloaded? like contract.