Suppose I have two tensors $A_{i_1,\ldots,i_M}$ and $B_{j_1,\ldots,j_N}$ where $M \neq N$ in general. We can define a tensor $C_{i_1,\ldots,i_M,j_1,\ldots,j_N}$ by

$$ C_{i_1,\ldots,i_M,j_1,\ldots,j_N} = A_{i_1,\ldots,i_M} + B_{j_1,\ldots,j_N} $$

I thought I could use torch.einsum to construct such a tensor in pytorch but it doesn't do the job unfortunately. Is there a way to do this in pytorch?

I wonder mostly if there's a pytorch API that can do this rather than me working out several tricks here and there.


1 Answer 1


I would tackle your problem as follows:

  1. Compute rank-$(M+N)$ tensors $\tilde{A}$ and $\tilde{B}$ that satisfy: $$ \forall j_1, ..., j_N\, \forall i_1, ..., i_M: \quad \tilde{A}_{i_1, ..., i_M, j_1, ..., j_N} = A_{i_1, ..., i_M}, \\ \forall i_1, ..., i_M\, \forall j_1, ..., j_N: \quad \tilde{B}_{i_1, ..., i_M, j_1, ..., j_N} = B_{j_1, ..., j_N}. $$

  2. Compute the answer by $$ C = \tilde{A} + \tilde{B}. $$ The second step is trivial and thus you just need to figure out how to do the first step. Computing the tensors $\tilde{A}$ and $\tilde{B}$ could possibly be done by expanding the dimensions to be $M+N$ and then stacking along the new dimensions or using torch.repeat_interleave. My answer is slightly different since you can also compute these larger tensors via outer products of $A$ and $B$ with tensors that contain only ones, i.e. $$ \tilde{A}_{i_1, ..., i_M, j_1, ..., j_N} = A_{i_1, ..., i_M} \times 1_{j_1, ..., j_N}, \\ \tilde{B}_{i_1, ..., i_M, j_1, ..., j_N} = 1_{i_1, ..., i_M} \times B_{j_1, ..., j_N}. $$ Here, is some PyTorch code that works for $M = 3$ and $N = 2$, but could also be modified for other values of $M$ and $N$:

    1. A_new = torch.einsum("ijk,lm->ijklm", A, torch.ones_like(B))

    2. B_new = torch.einsum("ijk,lm->ijklm", torch.ones_like(A), B)

    3. C = A_new + B_new

For more PyTorch API specific questions, you can also check out their official forum: https://discuss.pytorch.org/

  • 1
    $\begingroup$ The meaning of your equalities in point 1 is not clear. Can you please clarify your definitions? You have N+M indices on the lhs but only M or N on the rhs $\endgroup$ Commented Feb 23 at 22:47
  • $\begingroup$ @user8469759 Basically $\tilde{A}$ is just $A$ repeated along the new right $N$ dimensions and $\tilde{B}$ is just $B$ repeated along the new left $M$ dimensions. $\endgroup$ Commented Feb 24 at 14:22

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.