Any suggestions for multiplying a tensor (along an arbitrary axis) with a matrix in python?

Question

I have a tensor of rank N and bond dimension d with shape T=(d,...,d). I would like to multiply with a matrix M=(D,d) where D is not the same as d. The resulting tensor should have shape (d,...,d,D,d,...,d).

While I can do this to get eg (d,...,d,D,d,d) tensor:

np.einsum('ij, ...jkl->...ikl', M,T)

The tensor can be of quite a large rank and I need to do this several times. Therefore I want to avoid writing out each particular case as I did above as it would be impractical.

Can anyone suggest a better/more general/alternative way to do this? I would really appreciate any help. Thanks in advance.

Answer 1

Interesting problem. What you want is mutiply two tensors on specific dimensions and obtain a new tensor with a specific shape.

I tried different approaches and eventually came up with a composition of numpy.tensordot and numpy.rollaxis.

The former allows the tensor product between specified axes. The latter rotates the tensor to get the expected shape.

It was an interesting question, thanks. I hope I got this right, let me know.

import numpy as np
d=4
N=5
D=7
T = np.random.randint(0,9, (d,)*N)
M = np.random.randint(0,9, (D,d))

r = np.einsum('ij, ...jkl->...ikl', M,T)

i = 1
j = -3 

v = np.tensordot(M,T,axes=[[i],[j]])
v = np.rollaxis(v,0,j)

assert((v==r).all())