tfm.vision.box_ops.gather_instances

Gathers instances by indices.

selected_indices a Tensor of shape [batch, K] which indicates the selected indices in instance dimension (2nd dimension).
instances a Tensor of shape [batch, N, ...] where the 2nd dimension is the instance dimension to be selected from.
*aux_instances the additional Tensors whose shapes are in [batch, N, ...] which are the tensors to be selected from using the selected_indices.

selected_instances the tensor of shape [batch, K, ...] which corresponds to the selected instances of the instances tensor.
selected_aux_instances the additional tensors of shape [batch, K, ...] which corresponds to the selected instances of the aus_instances tensors.