tfr.utils.organize_valid_indices

Organizes indices in such a way that valid items appear first.

tfr.utils.organize_valid_indices(
    is_valid, shuffle=True, seed=None
)

is_valid A boolean Tensor for entry validity with shape [batch_size, list_size].
shuffle A boolean indicating whether valid items should be shuffled.
seed An int for random seed at the op level. It works together with the seed at global graph level together to determine the random number generation. See tf.set_random_seed.

A tensor of indices with shape [batch_size, list_size, 2]. The returned tensor can be used with tf.gather_nd and tf.scatter_nd to compose a new [batch_size, list_size] tensor. The values in the last dimension are the indices for an element in the input tensor.