I wasn't sure if I should post this on stackoverflow rather than here, but because I have to construct a specific tensor I think here is more suitable.

I have 2 tensors, $x \in \mathbb{R}^{M \times N \times L}$, and $d\in \mathbb{R}^{M \times L}$ my goal is to construct a new tensor defined as $(\cdot,\cdot) : \mathbb{R}^{M \times N \times L} \times \mathbb{R}^{M \times L} \to \mathbb{R}^{M\times N}$

defined as

$$ (x,d)_{ij} = \sum_{1 \leq l \leq L} x_{ijl}d_{il} $$

the closest thing I've found is the pytorch function tensordot

However given the following example

x = torch.tensor([[[ 1,  2,  3],[ 4,  5,  6],[ 7,  8,  9]],[[ 2,  4,  6],[ 8, 10, 12],[14, 16, 18]]])
d = torch.tensor([[1, 1, 1],[2, 2, 2]])
z = torch.tensordot(x,d,([2],[1]))

gives me the tensor

z =
tensor([[[ 6, 12],
         [15, 30],
         [24, 48]],
        [[12, 24],
         [30, 60],
         [48, 96]]])

Which corresponds to the formula

$$ z_{ijk} = \sum_{1 \leq l \leq L} x_{ijl}d_{kl} $$

to get the result I want I just have to take

$$ (x,d)_{ij} = z_{iji} $$

But I struggle to find an API that does this in pytorch, but I am sure there is one, I just don't know what to search for.

I've tried "trace" but I find only the trace of a square matrix. Does therefore anybody know if this operation has a name or how I can implement it using pytorch primitives without maybe resorting to for loops?


3 Answers 3


In your case, it might be convenient to perform the entire calculation with Einstein summation:

torch.einsum('ijl,il->ij', x, d)

This then directly computes the sum that you wrote down in your question and (up to transposition) yields the same solution as torch.diagonal(z, dim1=0, dim2=2) from @Vladimir Lysikov's answer.

  • $\begingroup$ Can this be used for tensor products? Suppose you have $A_{ij}$ and $B_{klm}$ two tensors. Then would we be getting with your call $C_{ijklm} = A_{ij}\times B_{klm}$? $\endgroup$ Feb 23 at 16:29
  • $\begingroup$ @user8469759 Yes, you could also use this for tensor products, e.g. torch.einsum('ijl,km->ijlkm', x, d) for the tensor product of x and d. (Alternatively, you could use torch.kron(x,d) but this stores the solution slightly differently then you want, e.g. the Kronecker product of two matrices is a very large matrix and not a rank-4 tensor.) $\endgroup$ Feb 23 at 17:48
  • $\begingroup$ More generally, the benefit of using torch.einsum is that you you can explicitly write down your tensor operations and don't need to write down the sums, since you automatically sum over repeated indices, as in my answer. The physics notation equivalent of torch.einsum('ijl,il->ij', x, d) is $c_{ij} = x_{ijl}d_{il}$, where the repetition of the index $l$ signifies that we mathematically have $c_{ij} = \sum_l x_{ijl}d_{il}$. $\endgroup$ Feb 23 at 17:52
  • $\begingroup$ Technically you have also a repetition of i as well in the example. But understand. Just out of curiosity is it also possible to construct, for example, 4D tensor from 2 2D tensors like $(a_{ij},b_{kl}) \to a_{ij} - b_{kl}$? $\endgroup$ Feb 23 at 19:17
  • $\begingroup$ I replied to your new question and that should hopefully resolve this comment $\endgroup$ Feb 23 at 20:55

I don't have PyTorch installed to check, but from the documentation it seems that torch.diagonal(z, dim1=0, dim2=2) should do the trick.


What I did at the end was something like

z = torch.tensordot(a,b,dim=[[2],[1]])
z = z[1:M,:,1:M]

where M is the common dimension size of the tensor z coming from the first line


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.