3

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.

ItsAmy
  • 844
  • 1
  • 6
  • 20
  • It might actually be better in terms of scalability and performance to make `f` work on tensors of dimensions `[a, n]` and `[b, n]` directly. – E_net4 Jun 22 '17 at 18:24
  • I guess I should clarify, I don't necessarily need `f` to be a separate function. I just want to apply _some_ operation that achieves the specified results. Having it operate on the entire tensors directly, as you said, would be the preferred solution. I'll clarify that now. – ItsAmy Jun 22 '17 at 18:27
  • 1
    I found this similar question, I'm looking into it now to see if I can expand it to functions other than euclidean distance: https://stackoverflow.com/questions/37009647/compute-pairwise-distance-in-a-batch-without-replicating-tensor-in-tensorflow?rq=1 – ItsAmy Jun 22 '17 at 18:46

1 Answers1

8
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>
mdaoust
  • 6,242
  • 3
  • 28
  • 29
  • Thanks, this was helpful! This works with Manhattan distance with very little modification, and I think I can make similar use of broadcasting for my other needs. – ItsAmy Jun 22 '17 at 22:40