Tensorflow efficient pairwise inner product

Question

In Tensorflow (python), given a matrix X of shape (n x d), where each row is a data point, I would like to compute the pairwise inner products of these n data points, i.e., the upper triangle of XX'.

Of course I could compute the whole XX' and fetch its upper triangle, but this means I would compute the off-diagonal elements twice. How to compute these efficiently in Tensorflow (python) by computing the inner product only once per pair?

Answer 1

With numpy, you can do this:

import numpy as np

A = np.random.randn(5, 3)
inds = np.triu_indices(5) # upper triangle indices

# expensive way to do it
ipu1 = np.dot(A, A.T)[inds]


# possibly less expensive way to do it.
ipu2 = np.einsum('ij,ij->i', A[inds[0]], A[inds[1]])

print(np.allclose(ipu1, ipu2))

This outputs True. Tensorflow does not have the triu_indices function build in, but it is not hard to write one if needed by looking at the numpy code. It does have einsum.