// CHECK: "tf.Iterator"
  // CHECK-SAME: output_shapes{{.*}}200x28x28x1{{.*}}200x10
  // CHECK-SAME: output_types = [f32, f32]
  // CHECK-SAME: shared_name = "_iterator1"
  %0 = "tf.AnonymousIteratorV3"() {output_shapes = [
    #tf_type.shape<200x28x28x1>,
    #tf_type.shape<200x10>], output_types = [f32, f32]} : () -> tensor<!tf_type.resource>
  // CHECK: "tf.Iterator"
  // CHECK-SAME: shared_name = "_iterator2"

  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
  SmallVector<mlir::Type, 2> input_types{input_type, indices_type};
  SmallVector<mlir::Type, 1> output_types{output_type};
  return createMaxUnpoolingFunc<2, 1>(builder, input_types, output_types);
}

template <int N>

    // CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "funcB"} : 1

    // CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "__inference_Dataset_flat_map_lambda_19"} : 1

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

           output_shapes = [#tf_type.shape<>], output_types = [i32],
           preserve_cardinality = false, sloppy = false,
           use_inter_op_parallelism = true,
           metadata = ""} : (tensor<*x!tf_type.variant>) -> tensor<!tf_type.variant>
  %5 = "tf.BatchDatasetV2"(%4, %0, %1) {device = "",
           output_shapes = [#tf_type.shape<>],
           output_types = [i32],
           parallel_copy = false,

      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "preserve_cardinality"
    value {
      b: false
    }
  }
  attr {

  computeInputsOutputs(ops, &inputs, &outputs);

  std::vector<Type> input_types;
  std::vector<Type> output_types;

  input_types.reserve(inputs.size());
  for (Value v : inputs) {
    input_types.push_back(v.getType());
  }
  output_types.reserve(outputs.size());
  for (Value v : outputs) {
    output_types.push_back(v.getType());
  }

  // Create the function.
  MLIRContext* context = ops[0]->getContext();

          }
          dim {
            size: 16
          }
        }
        shape {
          unknown_rank: true
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_BFLOAT16
        type: DT_FLOAT
        type: DT_DOUBLE
      }
    }
  }
}
node {
  name: "IteratorGetNextSync"

           output_shapes = [#tf_type.shape<>], output_types = [i32],
           preserve_cardinality = false, sloppy = false,
           use_inter_op_parallelism = true,
           metadata = ""} : (tensor<*x!tf_type.variant>, tensor<i32>) -> tensor<!tf_type.variant>
  %6 = "tf.BatchDatasetV2"(%5, %0, %1) {device = "",
           output_shapes = [#tf_type.shape<>],
           output_types = [i32],
           parallel_copy = false,