How to get a mask in the tensorflow graph by only tensor operations

Question

I have two tensors containing two groups of vectors d1 and d2. Both of d1 and d2 contain 5 two-dimensional vectors [d1 and d2 are changing in training loop]

import tensorflow as tf
import numpy as np
# random initialize
d1_np = np.random.rand(5,2) 
d2_np = np.random.rand(5,2) 

d1 = tf.Variable(initial_value = d1_np, dtype=tf.float32)
d2 = tf.Variable(initial_value = d2_np, dtype=tf.float32)

then I calculate the distance of them and getting a their cross distance matrix by cross_distance

dist_1_2 = cross_distance(d1, d2)  

so it produces a matrix of size 5x5 (the diagonal value is set to a very large value).

Then for each vector in d1, I got the index of its closest vector in d2 by

ind_min = tf.argmin(dist_1_2,axis=1)

ind_min gets values like [2 0 0 1 0] during runing

I then use tf.unique to get the unique index in ind_min

yv,idx = tf.unique(ind_min)

now yv becomes [2 0 1]. I want to set a mask and see whether the corresponding vectors in d2 is a closet vector to some vector in d1.

mask = tf.cast(tf.ones(5),tf.bool)

Now I hope to set the value of mask to zero for those index in yv. I tried:

mask[yv] = 0

('Tensor' object does not support item assignment) and

for ind in tf.unstack(yv):
    mask[yv] = 0

(Cannot infer num from shape (?,)) and it does not work.

The point is d1 and d2 is changing during some training process, so ind_min is not fixed but changing with the training loop.

Is there a way to get this mask dynamically?

Answer 1

Creating one-hot encoding of indices and adding along the first dimension should give you the mask. ie

mask = tf.reduce_sum(tf.one_hot(idx, 5), axis=0)

The mask size (hard-coded 5) can be replaced with d1.shape[0] .