%0:2 = "tf.TopKV2"(%arg0, %arg1): (tensor<16xf32>, tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>)
    func.return %0#0, %0#1: tensor<?xf32>, tensor<?xi32>
  }

// CHECK-LABEL:sin
// CHECK:  "tfl.sin"(%arg0) : (tensor<f32>) -> tensor<f32>
}

func.func @topk(%arg0: tensor<8xf32>, %arg1: tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>) {
  %0, %1 = "tf.TopKV2"(%arg0, %arg1) : (tensor<8xf32>, tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>)
  func.return %0, %1: tensor<?xf32>, tensor<?xi32>

// CHECK-LABEL: topk
// CHECK:  "tfl.topk_v2"(%arg0, %arg1)
// CHECK:  return
}

//
// values, indices = tfl.topkv2(%inputs, K)
// %1 = tfl.slice(values, 0, k)
// %2 = tfl.slice(indices,0, k)
// .... (values and indices only used for %1 and %2)
//
// %1 or %2 can be absent. If values and indices are only used here,
// this pattern can be replaced with (conceptually)
//
// %values, %indices = tfl.topkv2(%inputs, k)
// replace all use of %1 with values

// tf.TopKV2 legalization
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: topk_v2_non_const_k
func.func @topk_v2_non_const_k(%input: tensor<16xf32>, %k: tensor<i32>) -> (tensor<?xf32>, tensor<?xi32>) {
  // CHECK: tf.TopKV2

// CHECK:           %[[K:.*]] = "tf.Const"() <{value = dense<6> : tensor<i32>}> : () -> tensor<i32>
// CHECK:           %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) <{sorted = true}> : (tensor<3x6xf32>, tensor<i32>) -> (tensor<3x6xf32>, tensor<3x6xi32>)
// CHECK:           return %[[VALUES]], %[[INDICES]] : tensor<3x6xf32>, tensor<3x6xi32>
// CHECK:         }

            "_FusedBatchNormEx", "FusedBatchNormGrad", "FusedBatchNormGradV2",
            "FusedBatchNormGradV3"}},
          {"Conv", {"_FusedConv2D"}},
          {"SORT", {"TopKV2"}},  // XLA version much faster then TF version.
          {"MISC",
           // clang-format off
     {"ApproxTopK", "BroadcastTo", "ExpandDims", "Fill", "NoOp",

  %0:2 = "tf.TopKV2"(%input, %k) : (tensor<f32>, tensor<i32>) -> (tensor<*xf32>, tensor<*xi32>)
  func.return
}

// -----

func.func @testTopKV2WrongKRank(%input: tensor<8xf32>, %k: tensor<5xi32>) {
  // expected-error @+1 {{op requires k operand to be 0D tensor}}
  %0:2 = "tf.TopKV2"(%input, %k) : (tensor<8xf32>, tensor<5xi32>) -> (tensor<*xf32>, tensor<*xi32>)

                               /*outfeed_config=*/rewriter.getStringAttr(""));
    rewriter.eraseOp(op);
    return success();
  }
};

// Converts tf.TopKV2 to chlo.top_k.
class ConvertTopKV2Op : public OpRewritePattern<TF::TopKV2Op> {
 public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::TopKV2Op op,

    TF_Float32Tensor:$input,

    I64Attr:$k
  );

  let results = (outs
    TF_Float32Tensor:$topk,
    TF_Int32Tensor:$topk_indices
  );
}

def TF_TopKV2Op : TF_Op<"TopKV2", [Pure]> {
  let summary = [{
Finds values and indices of the `k` largest elements for the last dimension.
  }];

  let description = [{