// from `from` to `to`.
// TODO(b/158769932): This should be a general feature instead post some policy
// discussion.
inline void CopyDeviceAndUnderscoredAttributes(Operation *from, Operation *to) {
  auto device = mlir::StringAttr::get(from->getContext(), "device");
  CopyAttributes(from, to, [&device](const NamedAttribute &attr) {
    return attr.getName().strref().front() == '_' || attr.getName() == device;
  });
}

      resource_args.push_back(arg);
      has_resources = true;
    }
  }

  if (!in_order) {
    // Functions do not get reused in practice, so skip the check for if the
    // callee has been updated.
    StringAttr callee_sym = cluster_func_op.getFuncAttr().getAttr();
    MoveResourceArgsToEnd(symtab.lookup<func::FuncOp>(callee_sym));
  }
  builder.setInsertionPoint(cluster_func_op);
  auto xla_launch_op = builder.create<TF::XlaLaunchOp>(

    if (!backendConfig)
      return op->emitError(
          "custom_call has no 'composite.backend_config' attribute or the "
          "attribute is not a dictionary");

    auto name = mlir::dyn_cast<StringAttr>(backendConfig.get("name"));
    if (!name)
      return op->emitError(
          "backend_config has no 'name' key or the name value is not a string");

    auto attrs =

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

bool HasOutsideCompilationAttribute(Operation* op) {
  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {

SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;
  mlir::dyn_cast_or_null<StringAttr>(entry_function_attr.get("inputs"))
      .strref()
      .split(inputs, /*Separator=*/",");

  return inputs;
}

void ConvertMainArgAttrs(func::FuncOp main_func_op, const int arg_idx,

// If the given op pointer is a nullptr, returns false.
bool IsInStableHloOpRegion(Operation* op);

// Checks if a given einsum op is supported for XlaDotV2 quantization.
bool IsEinsumSupportedByXlaDotV2(StringAttr equation_attr);

// Gets the quantization method from `op`. It is retrieved from the
// `kQuantizationMethodAttr` string attribute. Returns
// `absl::InvalidArgumentError` when the attribute doesn't exist. Returns

  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());
    // Skip non-tpu device cluster_func.
    auto cluster_id =
        cluster_func_op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    llvm::SmallVector<int, 4> dynamic_shape_arg_index;

    // Traverse the operands of the cluster func op and find which operand

    // We may retrieve the original function's name from the attribute.
    // Functions without this attribute are ignored.
    auto original_func_name =
        func_op->getAttrOfType<StringAttr>("tf._original_func_name");
    if (original_func_name) {
      if (auto alias_itr = function_aliases.find(original_func_name.str());
          alias_itr != function_aliases.end()) {

  if (statistics_outputs.empty()) return failure();

  OpBuilder builder(&ctx);
  // Set the insertion point right before the return op.
  builder.setInsertionPoint(&region.back().back());

  StringAttr output_file_path_attr = builder.getStringAttr(output_file_path);
  ArrayAttr ids_attr = builder.getStrArrayAttr(ids);
  ArrayAttr calibration_methods_attr =
      builder.getI32ArrayAttr(calibration_methods);

    return absl::InternalError("Expected 1 argument for the main function.");
  }
  OpBuilder builder(module.getBodyRegion());
  // Create SessionInitializerOp; should reference main function.
  StringAttr func_name = main_func.getSymNameAttr();
  llvm::SmallVector<mlir::Attribute, 2> func_names = {
      mlir::SymbolRefAttr::get(builder.getContext(), func_name)};
  builder.create<tf_saved_model::SessionInitializerOp>(