TensorFlow Batch Outer Product

Question

I have the following two tensors:

x, with shape [U, N]
y, with shape [N, V]

I want to perform a batch outer product: I'd like to multiply each element in the first column of x by each element in the first row of y to get a tensor of shape [U, V], then the second column of x by the second row of y, and so on. The shape of the final tensor should be [N, U, V], where N is the batch size.

Is there any simple way to achieve this in TensorFlow? I've tried to use batch_matmul() without success.

Answer 1

Would the following work, using 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]

Answer 2

Perhaps, there is a more elegant solution using tf.einsum:

result = tf.einsum("un,nv->nuv", x, y)