}

      // Collect all the quantized outputs and replace them by the results of
      // the new quantized op.
      llvm::SmallDenseMap<Value, int> outputs_replaced;
      SmallVector<Type, 4> output_types;
      output_types.reserve(quantizing_op->getNumResults());
      for (const auto& enumerated_result :
           llvm::enumerate(quantizing_op->getResults())) {
        Value result = enumerated_result.value();

    SmallVector<int64_t, 3> slice_size = {1, num_rows, num_cols};
    Type slice_result_type = RankedTensorType::get(slice_size, element_type);
    llvm::SmallVector<Type, 4> output_types(batch_size, slice_result_type);
    auto split_op = rewriter.create<TF::SplitOp>(loc, output_types,
                                                 split_dimension_op.getOutput(),
                                                 reshape_op.getOutput());

              "max_percentile",
              rewriter.getF32FloatAttr(
                  calib_opts_.calibration_parameters().max_percentile())),
      };

      SmallVector<Type, 4> output_types{
          value.getType(),
          RankedTensorType::get({}, rewriter.getF32Type()),
          RankedTensorType::get({}, rewriter.getF32Type()),

  %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>

  %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

  for (Value v : inputs) operand_types.emplace_back(v.getType());
  llvm::SmallVector<Type, 4> output_types;
  output_types.reserve(outputs.size());
  for (Value v : outputs) output_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, output_types);

  FuncOp outlined_func =
      FuncOp::create(ops.front()->getLoc(), kHostFunctionAttr, func_type);

      }
      output_types = {new_output_type};
    } else {
      output_types = {output_type.clone(elem_type.getStorageType())};
    }

    SmallVector<Value> args = {q_op.getArg(), scale, zero_point};
    FlatSymbolRefAttr func_name =
        FlatSymbolRefAttr::get(rewriter.getStringAttr(kQuantizeFuncName));

    auto quantize_call = rewriter.create<TF::PartitionedCallOp>(
        loc, output_types, args, func_name,

    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

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