code: def top_kth_iterative(x, k):
    
    # The iterative computation is as follows:
    #
    # cur_x = x
    # for _ in range(k):
    #   top_x = maximum of elements of cur_x on the last axis
    #   cur_x = cur_x where cur_x < top_x and 0 everywhere else (top elements)
    #
    # We encode this computation in a TF graph using tf.foldl, so the inner
    # part of the above loop is called "next_x" and tf.foldl does the loop.
    def next_x(cur_x, _):
      top_x = tf.reduce_max(cur_x, axis=-1, keep_dims=True)
      return cur_x * to_float(cur_x < top_x)
    # We only do k-1 steps of the loop and compute the final max separately.
    fin_x = tf.foldl(next_x, tf.range(k - 1), initializer=tf.stop_gradient(x),
                     parallel_iterations=2, back_prop=False)
    return tf.stop_gradient(tf.reduce_max(fin_x, axis=-1, keep_dims=True))
docs: Compute the k-th top element of x on the last axis iteratively.
  
    This assumes values in x are non-negative, rescale if needed.
    It is often faster than tf.nn.top_k for small k, especially if k < 30.
    Note: this does not support back-propagation, it stops gradients!
  
    Args:
      x: a Tensor of non-negative numbers of type float.
      k: a python integer.
  
    Returns:
      a float tensor of the same shape as x but with 1 on the last axis
      that contains the k-th largest number in x.