How to Broadcast Sum Vector and Tensor?

Viewed 324

Suppose we have:

  • row vector V of shape (F,1), and
  • 4-D tensor T of shape (N, F, X, Y).

As a concrete example, let N, F, X, Y = 2, 3, 2, 2. Let V = [v0, v1,v2].

Then, I want to element-wise add v0 to the inner 2x2 matrix T[0,0], v1 to T[0,1], and v2 to T[0,2]. Similarly, I want to add v0 to T[1,0], v1 to T[1,1], and v2 to T[1,2].

So at the "innermost" level, the addition between the 2x2 matrix and a scalar, e.g. T[0,0] + v0, uses broadcasting to element-wise add v0. Then what I'm trying to do is apply that more generally to each inner 2x2.

I've tried using np.einsum() and np.tensordot(), but I couldn't figure out what each of those functions was actually doing on a more fundamental level, so I wanted to ask for a more step-by-step explanation of how this computation might be done.

Thanks

1 Answers

To multiply: You can simply translate your text into indices names of eisnum and it will take care of broadcasting:

TV = np.einsum('ijkl,j->ijkl',T,V)

To add: Simply add dimensions to your V using None to match up last two dimensions of T and broadcasting will take care of the rest:

TV = T + V[:,None,None]

Example input/output that shows the desired behavior of your output for adding:

T:

[[[[7 4]
   [5 9]]

  [[0 3]
   [2 6]]

  [[7 6]
   [1 1]]]


 [[[8 0]
   [8 7]]

  [[2 6]
   [9 2]]

  [[8 6]
   [4 9]]]]

V:

[0 1 2]

TV:

[[[[ 7  4]
   [ 5  9]]

  [[ 1  4]
   [ 3  7]]

  [[ 9  8]
   [ 3  3]]]


 [[[ 8  0]
   [ 8  7]]

  [[ 3  7]
   [10  3]]

  [[10  8]
   [ 6 11]]]]
Related