src/transformers/models/transfo_xl/modeling_tf_transfo_xl_utilities.py/TFAdaptiveSoftmaxMask/_gather_logprob
class TFAdaptiveSoftmaxMask:    def _gather_logprob(logprob, target):
        lp_size = shape_list(logprob)
        r = tf.range(lp_size[0], dtype=target.dtype)
        idx = tf.stack([r, target], 1)
        return tf.gather_nd(logprob, idx)