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

  std::vector<Type> input_types = GetValueTypes(inputs);
  std::vector<Type> output_types = GetValueTypes(outputs);
  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

    mlir::ArrayAttr dims = first_op.getPartitionDimsAttr();
    StringAttr sharding = first_op.get_XlaShardingAttr();
    Operation::result_type_range output_types = first_op.getResultTypes();
    result_op = builder.create<mlir::TF::TPUPartitionedOutputV2Op>(
        replicate_op.getLoc(), output_types, cluster.getResult(result_id), dims,
        sharding);

    results.insert(results.end(), result_op->getResults().begin(),

    const auto& arg = graph_op_def.output_arg(i);
    const auto& api_def_arg(api_def.out_arg(i));
    CHECK_EQ(arg.name(), api_def_arg.name());

    bool is_list = ArgIsList(arg);
    output_types.push_back(
        strings::StrCat("::tensorflow::", is_list ? "OutputList" : "Output"));
    output_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to()));
    is_list_output.push_back(is_list);
  }

    "tf.Yield"(%1) : (tensor<5xf32>) -> ()
  }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array<i32: 1, 2, 1>, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor<!tf_type.resource>, tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> {

  auto output = op->getResult(0);
  auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
  if (!output_type) return failure();

  // bias should be a vector sized of the last output dim.
  int64_t num_units = output_type.getDimSize(output_type.getRank() - 1);
  auto bias_type =
      mlir::RankedTensorType::get({num_units}, output_type.getElementType());

  mlir::DenseElementsAttr bias_attr;

      {"quantized_ops": ["${main_op}", "BiasAdd", "Relu"], "act_func": "internal_dequantize_and_relu_fn", "output_type": "f32"},
      {"quantized_ops": ["${main_op}", "BiasAdd", "Relu6"], "act_func": "internal_dequantize_and_relu6_fn", "output_type": "f32"},
    ]
    func.func @GenerateQuantizedFunctionName(${quantized_ops}, "${output_type}")(%input : tensor<*xi8>,
                           %filter : tensor<*xi8>, %bias : tensor<*xi32>,

        op->getLoc(), *output_type, op.getInput());
    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
        op,
        output_type->clone(
            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
                .getStorageType()),
        result);

    return success();
  }
};

// UniformDequantizeOp takes TF quantized types as input which would have been

      func.return %act : tensor<*x${output_type}>
    }

    parameters[
      {"quantized_ops": ["${main_op}"], "act_func": "internal_requantize_no_activation_fn", "output_type": "!tf_type.qint8"},
      {"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "!tf_type.qint8"},