CODEWITHSUNDEEP
pythontensorflowdeep-learning
piggs_boson
piggs_boson

Reputation: 1007

Pairwise sum in tensorflow

I have two tensors A and B, both of shape [10 5]. How do I compute a tensor C of shape [10 5 5] such that C[x, i, j] = A[x, i] + B[x, j]?

Edit: this is the sum analog of outer-product, not the outer product itself.

Upvotes: 4

Views: 711

Answers (3)

P-Gn
P-Gn

Reputation: 24581

You could rely on broadcasting.

op1 = tf.expand_dims(A, axis=2)
op2 = tf.expand_dims(B, axis=1)
C = tf.add(op1, op2)

Beware that the solution of @MaxB is not equivalent to this one as the operator [] is equivalent to a call to strided_slice, not expand_dims.

Upvotes: 1

MWB
MWB

Reputation: 12567

Slightly more readable and succinct than @user1735003's answer:

A[:, :, None] + B[:, None, :]

(Actually, the other answer swapped the axes)

Upvotes: 4

piggs_boson
piggs_boson

Reputation: 1007

Currently I'm using the property log(e^x * e^y) == x+y to accomplish addition from the matmul operation:

op1 = tf.reshape(tf.exp(A), [10, -1, 1])
op2 = tf.reshape(tf.exp(B), [10, 1, -1])
C = tf.log(tf.matmul(op1, op2))

but I'm guessing there would be a simpler/faster way out as well.

Upvotes: 0

Related Questions