TensorFlow Batch Outer Product
I have the following two tensors:
x, with shape [U, N]
y, with shape [N, V]
I want to do a batch external product: I would like to multiply every item in the first column x
by every item in the first row y
to get a tensor of the shape [U, V]
, then the second column by the x
second row y
, etc. The final tensor shape should be [N, U, V]
, where N
is the batch size.
Is there an easy way to achieve this in TensorFlow? I tried using batch_matmul()
without success.
source to share
Will the following work using it ? tf.batch_matmul()
print x.get_shape() # ==> [U, N]
print y.get_shape() # ==> [N, V]
x_transposed = tf.transpose(x)
print x_transposed.get_shape() # ==> [N, U]
x_transposed_as_matrix_batch = tf.expand_dims(x_transposed, 2)
print x_transposed_as_matrix_batch.get_shape() # ==> [N, U, 1]
y_as_matrix_batch = tf.expand_dims(y, 1)
print y_as_matrix_batch.get_shape() # ==> [N, 1, V]
result = tf.batch_matmul(x_transposed_as_matrix_batch, y_as_matrix_batch)
print result.get_shape() # ==> [N, U, V]
source to share