\\(predictions_i\\) be the predictions for all classes for example `i`,
  \\(targets_i\\) be the target class for example `i`,
  \\(out_i\\) be the output for example `i`,

$$out_i = predictions_{i, targets_i} \in TopKIncludingTies(predictions_i)$$
  }];

  let arguments = (ins
    Arg<TF_Float32Tensor, [{A `batch_size` x `classes` tensor.}]>:$predictions,
    Arg<TF_I32OrI64Tensor, [{A `batch_size` vector of class ids.}]>:$targets,