// 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>)

    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 = [{