func.func @tf2xla_fallback_op_approx_top_k(%arg0: tensor<16xf32>) -> (tensor<?xf32>, tensor<?xi32>) {
  %0:2 = "tf_device.cluster"() ({
    // CHECK: tf.ApproxTopK
    // CHECK-NOT: _xla_outside_compilation
    %1:2 = "tf.ApproxTopK"(%arg0) {k = 2} : (tensor<16xf32>) -> (tensor<?xf32>, tensor<?xi32>)
    tf_device.return %1#0, %1#1 : tensor<?xf32>, tensor<?xi32>

    if (op.getInputs().size() % 2 != 0) {
      return op.emitOpError() << "ApproxTopK takes an even number of operands.";
    }

    auto called_computations = op.getCalledComputations();
    if (called_computations.size() != 1) {
      return op.emitOpError()
             << "ApproxTopK takes exactly 1 called_computation.";
    }

  // CHECK:  return %2 : tensor<4x128x1xf32>
}

func.func @approx_top_k_with_max_k_last_reduction_dimension(%arg0: tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) {
  %values, %indices = "tf.ApproxTopK"(%arg0) {aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64} : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>)

    %0 = "tf.ReadVariableOp"(%arg0) {device = ""} : (tensor<!tf_type.resource<tensor<10x500xbf16>>>) -> tensor<10x500xbf16>
    // CHECK: mhlo.compare  GT

            "FusedBatchNormGradV3"}},
          {"Conv", {"_FusedConv2D"}},
          {"SORT", {"TopKV2"}},  // XLA version much faster then TF version.
          {"MISC",
           // clang-format off
     {"ApproxTopK", "BroadcastTo", "ExpandDims", "Fill", "NoOp",
      "Range", "Rank", "Reshape", "Shape", "ShapeN", "Size", "Squeeze",
      "Transpose", "ZerosLike", "OnesLike", "BiasAdd" /*PW + Broadcast*/,

// CHECK-SAME:                                        %[[ARG_3:.*]]: tensor<i32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) {
// CHECK:          %[[VALUES:.*]], %[[INDICES:.*]] = "tf.ApproxTopK"(%[[ARG_0]]) <{aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>)

    Res<TF_BoolTensor, [{The reduced tensor.}]>:$output
  );

  TF_DerivedOperandTypeAttr Tidx = TF_DerivedOperandTypeAttr<1>;

  let hasVerifier = 1;
}

def TF_ApproxTopKOp : TF_Op<"ApproxTopK", [Pure]> {
  let summary = [{
Returns min/max k values and their indices of the input operand in an approximate manner.
  }];

  let description = [{