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

[ifsheldon]

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.

[arogozhnikov]

looking at full problem, I'd implement it as named axes for state, and position-specified tensors for manipulations

x # axes: [q1, q2, q3, q4]
x = apply_gate(x, a, 'q1 q2 q3 -> q1 q2 q3') # a is 6 - dim, first three interpreted as input, last three as output
x = apply_gate(x, b, 'q3 q4 -> q3 q4')

Quite likely there are DSLs out there that do something similar

[ifsheldon]

I think your code is calculating the state after the two gates a and b, which is easier to implement, and we have quantum computing simulation toolkits for that. Mine was to calculate $b \otimes a$, not $(b \otimes a)x$.

[arogozhnikov]

that's correct.

I think you can just pass x = np.eye(2 ** 4) to compute the transform (assuming that code can work with batches of states). Or come up with additional syntax to do just that. If you preform such operations frequently, it may worth to cook (or find) specialized instrument.