code: def top_k_logits(logits, k):
      if k == 0:
          # no truncation
          return logits
  
      def _top_k():
          values, _ = tf.nn.top_k(logits, k=k)
          min_values = values[:, -1, tf.newaxis]
          return tf.compat.v1.where(
              logits < min_values,
              tf.ones_like(logits, dtype=logits.dtype) * -1e10,
              logits,
          )
  
      return tf.cond(
          pred=tf.equal(k, 0),
          true_fn=lambda: logits,
          false_fn=lambda: _top_k(),
      )
  
file_path: finetune/finetune/util/beam_search.py
imports: import math
  import numpy
  from finetune.util.shapes import shape_list
  import tensorflow
  from tensorflow.python.ops import inplace_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 for calling a global function with example usage: top_k_logits(logits, k) and returns: tf, logits, tf