code: def _my_top_k(x, k):
      """GPU-compatible version of top-k that works for very small constant k.
  
      Calls argmax repeatedly.
  
      tf.nn.top_k is implemented for GPU, but the gradient, sparse_to_dense,
      seems not to be, so if we use tf.nn.top_k, then both the top_k and its
      gradient go on cpu.  Once this is not an issue, this function becomes
      obsolete and should be replaced by tf.nn.top_k.
  
      Args:
        x: a 2d Tensor.
        k: a small integer.
  
      Returns:
        values: a Tensor of shape [batch_size, k]
        indices: a int32 Tensor of shape [batch_size, k]
      """
      if k > 10:
          return tf.nn.top_k(x, k)
      values = []
      indices = []
      depth = tf.shape(x)[1]
      for i in range(k):
          values.append(tf.reduce_max(x, 1))
          argmax = tf.argmax(x, 1)
          indices.append(argmax)
          if i + 1 < k:
              x += tf.one_hot(argmax, depth, -1e9)
      return tf.stack(values, axis=1), tf.to_int32(
          tf.stack(indices, axis=1)
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 11 for calling a global function with example usage: _my_top_k(x, k) and returns: tf,  tf, tf