finalize_func = @finalize,
      init_func = @init,
      next_func = @next,
      operandSegmentSizes = array<i32: 1, 2, 1>,
      output_shapes = [#tf_type.shape<>],
      output_types = [!tf_type.string],
      metadata = ""} : (
              tensor<4xf32>,
              tensor<3xf32>,
              tensor<!tf_type.resource>,
              tensor<2xf32>) -> tensor<!tf_type.variant>
  return

  %2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
  %3 = "tf.RangeDataset"(%0, %1, %2) {__op_key = 3: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
  // CHECK: tf_mlrt.executeop{{.*}}op: \22FlatMapDataset\22
  // CHECK-SAME: \22__inference_Dataset_flat_map_lambda_19\22

  replicas.reserve(num_replicas);
  auto devices = replicate_op.getDevices();

  // Collect result types and operands.
  Operation& terminator = replicate_op.GetBody().back();
  llvm::SmallVector<Type, 8> output_types(terminator.getOperandTypes());
  auto control_type = tf_executor::ControlType::get(island_op.getContext());
  llvm::SmallVector<Value, 8> replica_inputs(island_op.getControlInputs());

    mlir::ArrayAttr array_attr, llvm::SmallVector<mlir::Type> &output_types,
    llvm::SmallVector<mlir::DictionaryAttr> &output_attrs) {
  for (auto it : llvm::zip(
           types, array_attr.template getAsRange<mlir::DictionaryAttr>())) {
    if (mlir::isa<tfg::ControlType>(std::get<0>(it))) continue;
    output_types.push_back(std::get<0>(it));

    mlir::NamedAttrList list;

bool IsMustNotBeXlaOp(Operation* op) {
  for (auto& input : op->getOpOperands()) {
    if (TypeMustBeNonXLA(input.get().getType())) return true;
  }
  for (auto output_types : op->getResultTypes()) {
    if (TypeMustBeNonXLA(output_types)) return true;
  }
  return false;
}

// Check if the op must be compiled with XLA. If the op does not satisfy this

  }

  // CHECK-LABEL: @main
  func.func @main()
  attributes {tf_saved_model.exported_names = ["main"]} {
    %0 = "tf.Iterator"() {container = "", output_shapes = [#tf_type.shape<200x10>], output_types = [f32], shared_name = "foo_iterator"} : () -> tensor<!tf_type.resource>
    %1 = func.call @get_next(%0) : (tensor<!tf_type.resource>) -> tensor<200x10xf32>
    return
  }
}

// -----

        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :
         (tensor<!tf_type.string>, tensor<*x!tf_type.string>, tensor<!tf_type.string>, tensor<*xi64>, tensor<!tf_type.string>) -> tensor<*x!tf_type.variant>

void ParallelExecuteOp::build(OpBuilder& builder, OperationState& state,
                              int num_regions, TypeRange output_types) {
  DCHECK_GE(num_regions, 1);
  for (int i = 0; i < num_regions; ++i) {
    Region* region = state.addRegion();
    region->push_back(new Block);
  }
  state.addTypes(output_types);
}

Block& ParallelExecuteOp::GetRegionBlockWithIndex(unsigned index) {

        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :
         (tensor<!tf_type.string>, tensor<*x!tf_type.string>, tensor<!tf_type.string>, tensor<*xi64>, tensor<!tf_type.string>) -> tensor<*x!tf_type.variant>

  }

  auto anonymous_iterator = builder.create<AnonymousIteratorV3Op>(
      reduce_dataset.getLoc(),
      RankedTensorType::get({}, builder.getType<ResourceType>()),
      /*output_types=*/builder.getArrayAttr(type_attrs),
      /*shape_types=*/builder.getArrayAttr(shape_attrs));
  builder.create<MakeIteratorOp>(reduce_dataset.getLoc(),
                                 reduce_dataset.getInputDataset(),