nodes: 0, module; 1, function_definition; 2, function_name:_create_topk_unique; 3, parameters; 4, block; 5, identifier:inputs; 6, identifier:k; 7, expression_statement; 8, expression_statement; 9, expression_statement; 10, expression_statement; 11, expression_statement; 12, expression_statement; 13, expression_statement; 14, comment:# Select the current largest value k times and keep them in topk_r2. The; 15, comment:# selected largest values are marked as the smallest value to avoid being; 16, comment:# selected again.; 17, expression_statement; 18, expression_statement; 19, for_statement; 20, expression_statement; 21, expression_statement; 22, expression_statement; 23, expression_statement; 24, expression_statement; 25, expression_statement; 26, expression_statement; 27, return_statement; 28, comment:"""Creates the top k values in sorted order with indices.
  
    Args:
      inputs: A tensor with rank of 2. [batch_size, original_size].
      k: An integer, number of top elements to select.
  
    Returns:
      topk_r2: A tensor, the k largest elements. [batch_size, k].
      topk_indices_r2: A tensor, indices of the top k values. [batch_size, k].
    """; 29, assignment; 30, assignment; 31, assignment; 32, assignment; 33, assignment; 34, assignment; 35, assignment; 36, assignment; 37, identifier:i; 38, call; 39, block; 40, assignment; 41, assignment; 42, assignment; 43, assignment; 44, assignment; 45, assignment; 46, assignment; 47, expression_list; 48, identifier:height; 49, subscript; 50, identifier:width; 51, subscript; 52, identifier:neg_inf_r0; 53, call; 54, identifier:ones; 55, call; 56, identifier:neg_inf_r2; 57, binary_operator:ones * neg_inf_r0; 58, identifier:inputs; 59, call; 60, identifier:tmp; 61, identifier:inputs; 62, identifier:topk_r2; 63, call; 64, identifier:range; 65, argument_list; 66, expression_statement; 67, expression_statement; 68, expression_statement; 69, expression_statement; 70, expression_statement; 71, identifier:log2_ceiling; 72, call; 73, identifier:next_power_of_two; 74, binary_operator:1 << log2_ceiling; 75, identifier:count_mask; 76, binary_operator:next_power_of_two - 1; 77, identifier:mask_r0; 78, call; 79, identifier:mask_r2; 80, call; 81, identifier:topk_r2_s32; 82, call; 83, identifier:topk_indices_r2; 84, call; 85, identifier:topk_r2; 86, identifier:topk_indices_r2; 87, attribute; 88, integer:0; 89, attribute; 90, integer:1; 91, attribute; 92, argument_list; 93, attribute; 94, argument_list; 95, identifier:ones; 96, identifier:neg_inf_r0; 97, attribute; 98, argument_list; 99, attribute; 100, argument_list; 101, identifier:k; 102, assignment; 103, assignment; 104, assignment; 105, assignment; 106, assignment; 107, identifier:int; 108, argument_list; 109, integer:1; 110, identifier:log2_ceiling; 111, identifier:next_power_of_two; 112, integer:1; 113, attribute; 114, argument_list; 115, attribute; 116, argument_list; 117, attribute; 118, argument_list; 119, attribute; 120, argument_list; 121, identifier:inputs; 122, identifier:shape; 123, identifier:inputs; 124, identifier:shape; 125, identifier:tf; 126, identifier:constant; 127, unary_operator; 128, keyword_argument; 129, identifier:tf; 130, identifier:ones; 131, list; 132, keyword_argument; 133, identifier:tf; 134, identifier:where; 135, call; 136, identifier:neg_inf_r2; 137, identifier:inputs; 138, identifier:tf; 139, identifier:zeros; 140, list; 141, keyword_argument; 142, identifier:kth_order_statistic; 143, call; 144, identifier:k_mask; 145, call; 146, identifier:topk_r2; 147, call; 148, identifier:ge_r2; 149, call; 150, identifier:tmp; 151, call; 152, call; 153, identifier:tf; 154, identifier:constant; 155, identifier:count_mask; 156, identifier:tf; 157, identifier:fill; 158, list; 159, identifier:mask_r0; 160, identifier:tf; 161, identifier:bitcast; 162, identifier:topk_r2; 163, attribute; 164, attribute; 165, identifier:bitwise_and; 166, identifier:topk_r2_s32; 167, identifier:mask_r2; 168, attribute; 169, identifier:dtype; 170, attribute; 171, identifier:height; 172, identifier:width; 173, identifier:dtype; 174, attribute; 175, attribute; 176, argument_list; 177, identifier:height; 178, identifier:k; 179, identifier:dtype; 180, attribute; 181, attribute; 182, argument_list; 183, attribute; 184, argument_list; 185, attribute; 186, argument_list; 187, attribute; 188, argument_list; 189, attribute; 190, argument_list; 191, attribute; 192, argument_list; 193, identifier:height; 194, identifier:k; 195, identifier:tf; 196, identifier:int32; 197, identifier:tf; 198, identifier:bitwise; 199, identifier:np; 200, identifier:inf; 201, identifier:tf; 202, identifier:float32; 203, identifier:tf; 204, identifier:float32; 205, identifier:tf; 206, identifier:is_nan; 207, identifier:inputs; 208, identifier:tf; 209, identifier:float32; 210, identifier:tf; 211, identifier:reduce_max; 212, identifier:tmp; 213, keyword_argument; 214, keyword_argument; 215, identifier:tf; 216, identifier:tile; 217, call; 218, list; 219, identifier:tf; 220, identifier:where; 221, identifier:k_mask; 222, call; 223, identifier:topk_r2; 224, identifier:tf; 225, identifier:greater_equal; 226, identifier:inputs; 227, call; 228, identifier:tf; 229, identifier:where; 230, identifier:ge_r2; 231, identifier:neg_inf_r2; 232, identifier:inputs; 233, identifier:math; 234, identifier:ceil; 235, call; 236, identifier:axis; 237, integer:1; 238, identifier:keepdims; 239, True; 240, attribute; 241, argument_list; 242, identifier:height; 243, integer:1; 244, attribute; 245, argument_list; 246, attribute; 247, argument_list; 248, attribute; 249, argument_list; 250, identifier:tf; 251, identifier:expand_dims; 252, call; 253, integer:0; 254, identifier:tf; 255, identifier:tile; 256, identifier:kth_order_statistic; 257, list; 258, identifier:tf; 259, identifier:tile; 260, identifier:kth_order_statistic; 261, list; 262, identifier:math; 263, identifier:log; 264, call; 265, integer:2; 266, attribute; 267, argument_list; 268, integer:1; 269, identifier:k; 270, integer:1; 271, identifier:width; 272, identifier:float; 273, argument_list; 274, identifier:tf; 275, identifier:equal; 276, call; 277, call; 278, call; 279, attribute; 280, argument_list; 281, attribute; 282, argument_list; 283, identifier:int; 284, argument_list; 285, identifier:tf; 286, identifier:range; 287, identifier:k; 288, identifier:tf; 289, identifier:fill; 290, list; 291, identifier:i; 292, identifier:width; 293, identifier:k
code: def _create_topk_unique(inputs, k):
    """Creates the top k values in sorted order with indices.
  
    Args:
      inputs: A tensor with rank of 2. [batch_size, original_size].
      k: An integer, number of top elements to select.
  
    Returns:
      topk_r2: A tensor, the k largest elements. [batch_size, k].
      topk_indices_r2: A tensor, indices of the top k values. [batch_size, 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)
  
    # Select the current largest value k times and keep them in topk_r2. The
    # selected largest values are marked as the smallest value to avoid being
    # selected again.
    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