code: 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)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/transfo_xl/modeling_tf_transfo_xl_utilities.py
imports: import tensorflow
  from ...tf_utils import shape_list
label_desc: Use image node_id 4 for calling the TFAdaptiveSoftmaxMask obj's underlying member method code with example usage: obj._gather_logprob(logprob, target) and returns: tf