code: def top_k_softmax(x, k):
      """Calculate softmax(x), select top-k and rescale to sum to 1.
  
      Args:
        x: Input to softmax over.
        k: Number of top-k to select.
  
      Returns:
        softmax(x) and maximum item.
      """
      x = tf.nn.softmax(x)
      top_x, _ = tf.nn.top_k(x, k=k + 1)
      min_top = tf.reduce_min(top_x, axis=-1, keep_dims=True)
      x = tf.nn.relu((x - min_top) + 1e-12)
      x /= tf.reduce_sum(x, axis=-1, keep_dims=True)
      return x, tf.reduce_max(top_x, axis=-1)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from functools import partial
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
  from tensorflow.python.training import moving_averages
label_desc: Use image node_id 10 for calling a global function with example usage: top_k_softmax(x, k) and returns: x,  tf