code: def _top_k_logits(logits, k):
      """Adapted from
      https://github.com/openai/gpt-2/blob/master/src/sample.py#L63-L77
      """
      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.where(
              logits < min_values,
              tf.ones_like(logits, dtype=logits.dtype) * -1e10,
              logits,
          )
  
      return tf.cond(
          tf.equal(k, 0),
          lambda: logits,
          lambda: _top_k(),
      )
  
imports: import tensorflow
  from tensorflow.python.ops import array_ops
  from tensorflow_probability import distributions
  from texar.tf.modules.decoders.tf_helpers import Helper, TrainingHelper, GreedyEmbeddingHelper
  from texar.tf.modules.embedders.embedder_utils import soft_embedding_lookup
  from texar.tf.utils import utils
label_desc: Use image node_id 5 for calling a global function with example usage: _top_k_logits(logits, k) and returns: tf, logits, tf