如何结合使用 tf.gather 和 tf.where

Question

The following is not working due to the shape of tf.where() . Is there a nice way to fix this?

I want the values of tensor_y where tensor_x fulfills a condition (eg == value ). Important, the tensors have batch_dims = 1 .

tensor_x = tf.random.uniform(shape=[2, 10], minval=-1, maxval=2, dtype=tf.int32)
tensor_y = tf.random.uniform(shape=[2, 10], minval=-1, maxval=2, dtype=tf.int32)
trues = tf.where(tensor_y ==1)


new_tensor = tf.gather(tensor_y, axis=-1, indices = trues,batch_dims=1)

What I am doing now works*, but it is not so efficient I think:

new_tensor = tf.stack([tf.gather(tensor_y[i,:], tf.where(tensor_x[i,:] == 1)) for i in range(tensor_x.shape[0])])

*sometimes (I don't know under which conditions) I get his error:

Shapes of all inputs must match: values[0].shape = [3,1].= values[1],shape = [6:1] [Op:Pack] name: stack

Answer 1

Is this what you need?

tensor_x = tf.random.uniform(shape=[2, 10], minval=-1, maxval=2, dtype=tf.int32)
tensor_y = tf.random.uniform(shape=[2, 10], minval=-1, maxval=2, dtype=tf.int32)

new_tensor = tensor_y[(tensor_x==1)]