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(),
      )
  
imports: import tensorflow
  from gpt_2_simple.src import model
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