code: def _sample_next_tokens(logits):
      """Sample tokens for next timestep."""
      batch_size = common_layers.shape_list(logits)[0]
      next_tokens = tf.random.categorical(logits, 1)
      next_tokens = tf.cast(next_tokens, tf.int32)
      next_tokens = tf.reshape(next_tokens, [batch_size])
      return next_tokens
  
file_path: tensor2tensor/tensor2tensor/utils/scheduled_sampling.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.python.ops import inplace_ops
label_desc: Use image node_id 4 for calling a global function with example usage: _sample_next_tokens(logits) and returns: next_tokens