General use of ellipsis in np.einsum

Question

I am trying to implement a tensor network kind of calculation using np.einsum. The goal of the calculation is to have a local 2x2 matrix act sequentially on a tensor of size (2,2,2,...,2,2,2) where the number of dimensions can be arbitrary (as long as memory allows). I have a working naive implementation for a tensor of size (2,2,2)

X = np.array([[1,2],[3,4]])
Y = np.ones(2**3).reshape((2,2,2))
Aijk = np.einsum('zi,zjk', X, Y)
Bijk = np.einsum('zi,jzk', X, Aijk)
Cijk = np.einsum('zi,jkz', X, Bijk)

As can be seen, the index to be summed over should move to the right for each subsequent calculation. I could generate the appropriate explicit string but I'm wondering if it possible to implement this using ellipses but I am not sure how to specify the index position if the index is not the left- or rightmost index.

X = np.array([[1,2],[3,4]])
Y = np.ones(2**n).reshape(shape) # where shape is the appropriate n-tuple
A = np.einsum('a...,a...', X, Y)
B = np.einsum('b...,...b...', X, A) # how do I specify that the summed index of A is the second one?
...
N = np.einsum('l...,...l', X, M)

I hope it is clear what I'm trying to achieve.

Answer 1

Let's take small steps to add ellipsis to your example.

In [3]: X = np.array([[1,2],[3,4]])
   ...: Y = np.ones(2**3).reshape((2,2,2))
   ...: Aijk = np.einsum('zi,zjk', X, Y)

In [4]: Aijk
Out[4]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

Making the RHS explicit:

In [5]: np.einsum('zi,zjk->ijk', X, Y)
Out[5]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

Replace the repeated jk with ellipsis:

In [6]: np.einsum('zi,z...->i...', X, Y)
Out[6]: 
array([[[4., 4.],
        [4., 4.]],

       [[6., 6.],
        [6., 6.]]])

And trying that with other Y shapes (always keeping first dimension, 2 to match on z:

In [7]: np.einsum('zi,z...->i...', X, np.ones((2,2,2,2)))
Out[7]: 
array([[[[4., 4.],
         [4., 4.]],

        [[4., 4.],
         [4., 4.]]],


       [[[6., 6.],
         [6., 6.]],

        [[6., 6.],
         [6., 6.]]]])

In [12]: np.einsum('zi,z...->i...', X, np.ones((2,3)))
Out[12]: 
array([[4., 4., 4.],
       [6., 6., 6.]])

In [13]: np.einsum('zi,z...->i...', X, np.ones((2,1,2,3)))
Out[13]: 
array([[[[4., 4., 4.],
         [4., 4., 4.]]],


       [[[6., 6., 6.],
         [6., 6., 6.]]]])

I think I could use 'zi,z...->...i', 'zi,jz...->ij...', 'zi,jz...->i...j'. (but I haven't tested them). But I'm pretty sure I can't use '...i...'. And in general can't use ellispis to specify the relative location of the summing 'z' dimension.