// CHECK-LABEL: func @skip_noncompiled_reduce_dataset
func.func @skip_noncompiled_reduce_dataset(
      %arg0 : tensor<!tf_type.variant>,
      %arg1: tensor<i64>
    ) {
    // CHECK: tf.ReduceDataset
    %1 = "tf.ReduceDataset"(%arg0, %arg1) {
      Targuments = [],
      Tstate = [i64], device = "",
      f = @__reduce_func_0, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],

      output_shapes = [#tf_type.shape<?>],
      output_types = [!tf_type.string]} : (tensor<*x!tf_type.variant>, tensor<i64>) -> tensor<!tf_type.variant>

    %1 = "tf.ReduceDataset"(%0, %cst_1, %arg0) {
      Targuments = [!tf_type.resource],
      Tstate = [i32], device = "",
      f = @__reduce_func, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],

    Location loc = reduce_dataset.getLoc();

    // Get reduce function signature for dataset iteration types.
    // Note: lookupSymbol is a linear lookup which means the overall
    // complexity = # ReduceDataset ops x # of functions in module.
    func::FuncOp reduce_func =
        function->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(
            reduce_dataset.getF());

  let summary = "Decomposes ReduceDataset op into dataset operations.";

  let description = [{
    Decomposes ReduceDataset op into a while loop that iterates the dataset and calls
    into the reduction function.  This decomposition is only done if the
    ReduceDataset op is marked for compilation with the _xla_compile_device_type
    attribute.

### `-tf-decompose-reduce-dataset`

_Decomposes ReduceDataset op into dataset operations._

  Decomposes ReduceDataset op into a while loop that iterates the dataset and calls
  into the reduction function.  This decomposition is only done if the
  ReduceDataset op is marked for compilation with the _xla_compile_device_type
  attribute.

  For example, for the following function the ReduceDataset op:

  ```mlir

// Performs specific fusion for GPU targets.
std::unique_ptr<OperationPass<func::FuncOp>> CreateGpuOpFusionPass();

// Creates a pass that decomposes to be compiled ReduceDataset ops into a while
// loop that iterates the dataset and calls the reduction function.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeReduceDatasetPass();

      output_types = [!tf_type.string, f32, f32, i32, i32, f32, f32, !tf_type.string, !tf_type.string],
      metadata = ""} : (tensor<*x!tf_type.variant>, tensor<i64>) -> tensor<!tf_type.variant>

    %117:6 = "tf.ReduceDataset"(%51, %cst_12, %cst_12, %cst_12, %cst_12, %cst_13, %cst_13, %arg2, %arg3, %arg4, %arg5, %arg6, %arg6, %arg6, %arg6, %arg7, %arg8, %arg9, %arg10, %cst_0, %arg0, %arg6, %arg11, %arg6, %arg6, %arg6, %arg6, %arg9) {

              "Inferring shape for ReduceDataset with #states = "
                  << num_states << " , #input_elements = " << num_input_elements
                  << " , and #captured_arguments = " << num_captured_arguments);
  if (num_states + num_input_elements + num_captured_arguments !=
      f.getNumArguments()) {
    op->emitOpError(
        "propagating shapes for ReduceDataset failed due to inconsistent "

}

//===----------------------------------------------------------------------===//
// tf.data ops
//===----------------------------------------------------------------------===//

def TF_ReduceDatasetOp : TF_Op<"ReduceDataset", [SameVariadicOperandSize]> {
  let summary = [{
    Reduces the input dataset to a singleton using a reduce function.
  }];

  let arguments = (ins
    TF_VariantTensor:$input_dataset,