}

// CHECK-LABEL: iterators
func.func private @iterators(%arg0 : tensor<!tf_type.variant>) {
  %0 = "tf.Iterator"() {container = "", output_shapes = [#tf_type.shape<200x28x28x1>, #tf_type.shape<200x10>], output_types = [f32, f32], shared_name = "_iterator1"} : () -> tensor<!tf_type.resource>
  %1 = func.call @id(%arg0) : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>

  }

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

        // expected-remark@above {{ID: 0}}

    %0 = "tf.TakeDataset"(%arg0, %cst) {device = "", metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<!tf_type.variant>, tensor<i64>) -> tensor<!tf_type.variant>

  bool OpProducesOrConsumesVariant(const Node& node) const {
    auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; };
    return absl::c_any_of(node.input_types(), is_variant) ||
           absl::c_any_of(node.output_types(), is_variant);
  }

  bool HasXLAKernel(const Node& node,
                    string* uncompilable_reason = nullptr) const;

  static void MaybeMarkUncompilableNode(
      const absl::string_view reason,

    // to the region block's return.
    bool RegionWrapsSingleOp(unsigned index);
  }];

  let builders = [
    OpBuilder<(ins "int":$num_regions, "TypeRange":$output_types)>,
  ];

  let hasVerifier = 1;
}

def TfDevice_ReplicateOp : TfDevice_Op<"replicate",
    [SingleBlockImplicitTerminator<"ReturnOp">, AttrSizedOperandSegments]> {