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.
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]],