(inputs, k)
return topk_r2, topk_indices_r2
Creates the top k values in sorted order with indices.
Creates the top k values in sorted order with indices.
Creates the top k values in sorted order with indices .
def create_topk_unique(inputs, k):
  """Creates the top k values in sorted order with indices."""
  height = inputs.shape[0]
  width = inputs.shape[1]
  neg_inf_r0 = tf.constant(-np.inf, dtype=tf.float32)
  ones = tf.ones([height, width], dtype=tf.float32)
  neg_inf_r2 = ones * neg_inf_r0
  inputs = tf.where(tf.is_nan(inputs), neg_inf_r2, inputs)

  tmp = inputs
  topk_r2 = tf.zeros([height, k], dtype=tf.float32)
  for i in range(k):
    kth_order_statistic = tf.reduce_max(tmp, axis=1, keepdims=True)
    k_mask = tf.tile(
        tf.expand_dims(tf.equal(tf.range(k), tf.fill([k], i)), 0), [height, 1])
    topk_r2 = tf.where(k_mask, tf.tile(kth_order_statistic, [1, k]), topk_r2)
    ge_r2 = tf.greater_equal(inputs, tf.tile(kth_order_statistic, [1, width]))
    tmp = tf.where(ge_r2, neg_inf_r2, inputs)

  log2_ceiling = int(math.ceil(math.log(float(int(width)), 2)))
  next_power_of_two = 1 << log2_ceiling
  count_mask = next_power_of_two - 1
  mask_r0 = tf.constant(count_mask)
  mask_r2 = tf.fill([height, k], mask_r0)
  topk_r2_s32 = tf.bitcast(topk_r2, tf.int32)
  topk_indices_r2 = tf.bitwise.bitwise_and(topk_r2_s32, mask_r2)
  return topk_r2, topk_indices_r2
def create_topk_unique ( inputs , k ) : height = inputs . shape [ 0 ] width = inputs . shape [ 1 ] neg_inf_r0 = tf . constant ( - np . inf , dtype = tf . float32 ) ones = tf . ones ( [ height , width ] , dtype = tf . float32 ) neg_inf_r2 = ones * neg_inf_r0 inputs = tf . where ( tf . is_nan ( inputs ) , neg_inf_r2 , inputs ) tmp = inputs topk_r2 = tf . zeros ( [ height , k ] , dtype = tf . float32 ) for i in range ( k ) : kth_order_statistic = tf . reduce_max ( tmp , axis = 1 , keepdims = True ) k_mask = tf . tile ( tf . expand_dims ( tf . equal ( tf . range ( k ) , tf . fill ( [ k ] , i ) ) , 0 ) , [ height , 1 ] ) topk_r2 = tf . where ( k_mask , tf . tile ( kth_order_statistic , [ 1 , k ] ) , topk_r2 ) ge_r2 = tf . greater_equal ( inputs , tf . tile ( kth_order_statistic , [ 1 , width ] ) ) tmp = tf . where ( ge_r2 , neg_inf_r2 , inputs ) log2_ceiling = int ( math . ceil ( math . log ( float ( int ( width ) ) , 2 ) ) ) next_power_of_two = 1 << log2_ceiling count_mask = next_power_of_two - 1 mask_r0 = tf . constant ( count_mask ) mask_r2 = tf . fill ( [ height , k ] , mask_r0 ) topk_r2_s32 = tf . bitcast ( topk_r2 , tf . int32 ) topk_indices_r2 = tf . bitwise . bitwise_and ( topk_r2_s32 , mask_r2 ) return topk_r2 , topk_indices_r2