CODEWITHSUNDEEP
pythontensorflowlinear-algebra
ItsAmy
ItsAmy

Reputation: 844

Apply function for every pair of elements in two Tensors in Tensorflow

I have two rank-2 Tensors with equal sizes along the second dimension, but unequal along the first. For example, tensor A of shape [a, n] and tensor B of shape [b, n]. They can be regarded as two arrays containing vectors of length n.

I have a function f which takes two inputs, each a tensor of shape [n], and returns a scalar. I want to apply this function to each pair of vectors in A and B with the result being a tensor C of shape [a, b] such that, for each location (i, j) in C, C[i, j] = f(A[i], B[j]).

If these were just regular Numpy arrays, I could accomplish this with the following code:

# Assume a, b, and n are integers, and A and B are Numpy arrays
C = numpy.zeros((a, b))
for i in range(0, a):
    for j in range(0, b):
        C[i, j] = f(A[i], B[j])
return C

If this could be accomplished in such a way that f simply takes A and B as input and returns C, that would be the preferred solution, so that everything happens as proper tensor operations, so that it can all be properly parallelized by Tensorflow. Just so long as the end result is the same.

I have found a solution to this problem specifically for when f calculates the euclidean distance between each pair of vectors. I would like to extend this to other functions, such as cosine distance or Manhattan (L1) distance.

Upvotes: 3

Views: 2513

Answers (1)

mdaoust
mdaoust

Reputation: 6367

a = tf.random_normal([10,5])
b = tf.random_normal([20,5])

I would start by re-orienting the two arrays like this:

a = a[:,tf.newaxis,:]
b = b[tf.newaxis,:,:]

Now the shapes are [a,1,n] and [1,b,n], so we can broad-cast a subtraction to calculate the delta for each pair:

delta = (a-b)

This has a shape of [a,b,n].

Now the Euclidean distance is straight forward. (axis=-1 summs over the last axis):

distance = tf.reduce_sum(delta**2,axis = -1)**0.5

And you're done:

print(distance)
<tf.Tensor 'pow_3:0' shape=(10, 20) dtype=float32>

Upvotes: 8

Related Questions