Tensorflow - constructing a tensor from particular values extracted from two different tensors

Question

I'm trying to construct a single tensor using values from two different tensors and an array of two dimensional indices, in a manner compatible with TensorFlow autodiff.

In a first step I want to extract the elements of a tensor D of shape (n,n) whose values are the same as those in another tensor a. In particular, I'm looking for a better way to implement the following loop:

a = []
for i in range(len(f)):
    a.append(tf.where(tf.experimental.numpy.isclose(f[I], D, atol=1e-6))[0])
P_x = tf.gather(D,a)

In the append step, I'm just using the first instance where the values are equal because the function I'm interested in is independent of this choice. I need to use isclose because the two arrays are float32 arrays and are not exactly equal to one another.

Then in a second step I want to combine P_x with P_y = tf.gather(g, indices) to construct a tensor P. Assume that P_x and P_y are both of shape (n, ). Then,

P = [[P_x[0], P_y[0]],[P_x[1], P_y[1]], ..., [P_x[n], P_y[n]] ]

I'm pretty new to TensorFlow, so despite looking through the docs I don't see a way to do all of these operations using gather, scatter etc., which seems to be necessary to make autodiff work. When I use loops and other methods, I get gradients = none.

Answer 1

For the first step, you can reduce the loop into matrix operation by finding the closest match using broadcasting.

indices = tf.reduce_sum(tf.math.abs(D[:,None] - a), 2)     
#numpy is_close

tf.gather(D,tf.where(indices < 1e-6)[:,0])

Example:

D = tf.random.normal(shape=(5,3))
a = tf.concat([tf.random.normal(shape=(2,3)), D[:2],], axis=0)

#Here a last 2 rows of `a` are made same as first two rows of D.
#D is array([[ 0.6221494 ,  0.39071774,  0.5728211 ],
   [ 0.926828  ,  0.8460992 ,  0.08634651],
   [-0.39511812, -0.02012417,  1.0490925 ],
   [-0.31207308,  0.41652176,  0.85152763],
   [-1.27271   , -0.09542792, -0.16090107]]
#a is array([[ 0.9826471 ,  0.25055575, -0.4920534 ],
   [-0.3222343 ,  0.91883016,  1.2904693 ],
   [ 0.6221494 ,  0.39071774,  0.5728211 ],
   [ 0.926828  ,  0.8460992 ,  0.08634651]]

numpy.is_close() operation

indices = tf.reduce_sum(tf.math.abs(D[:,None] - a), 2)  
#this compares each row of D with each row of a. So we get a (5,4) matrix for the above example.

Gather D close to a:

tf.gather(D,tf.where(indices < 1e-6)[:,0])
#output
array([[0.6221494 , 0.39071774, 0.5728211 ],
       [0.926828  , 0.8460992 , 0.08634651]],