OptionalOpShardingVector& sharding_for_args) {
  // Look up function definition from module.
  Block& function_block = func.front();

  sharding_for_args.reserve(function_block.getNumArguments());

  // Iterate through operands of `cluster_func`.
  // The computation operand can either be:
  //   1) a TPUPartitionedInput Op if the input has a non-resource type;
  //   2) a ReadVariableOp else.
  //

    // No need to create a mapping when there is no usage - it will not affect
    // the cloning.
    if (src_arg.use_empty()) continue;

    const unsigned main_arg_idx = main_func_op.getNumArguments();

    const DictionaryAttr main_arg_attr =
        src_func_op.getArgAttrDict(src_arg_idx);

    main_func_op.insertArgument(main_arg_idx, src_arg.getType(), main_arg_attr,

  constexpr char kAliasingAttr[] = "tf.aliasing_output";
  llvm::SmallDenseMap<unsigned, unsigned> output_to_input_alias;
  unsigned num_arguments = main_fn.getNumArguments();
  for (unsigned arg_index = 0; arg_index < num_arguments; ++arg_index) {
    if (auto aliasing_output = main_fn.getArgAttrOfType<mlir::IntegerAttr>(
            arg_index, kAliasingAttr))

      unique_output_names.push_back(
          exporter.op_to_name_.GetUniqueName(tensor_id_node));
    }
  }

  if (!input_names.empty()) {
    TF_RET_CHECK(input_names.size() == block.getNumArguments());
    for (const auto& it : llvm::enumerate(function.getArguments())) {
      // TODO(lyandy): Update when changing feed/fetch import.
      std::string name(input_names[it.index()]);

      unique_output_names.push_back(
          exporter.op_to_name_.GetUniqueName(tensor_id_node));
    }
  }

  if (!input_names.empty()) {
    TF_RET_CHECK(input_names.size() == block.getNumArguments());
    for (const auto& it : llvm::enumerate(function.getArguments())) {
      // TODO(lyandy): Update when changing feed/fetch import.
      std::string name(input_names[it.index()]);

void RemoveQuantizationAdaptorOps(func::FuncOp func) {
  mlir::OpBuilder builder(func.getBody());
  auto& bb = func.front();
  auto loc = func.getLoc();

  int num_args = bb.getNumArguments();
  llvm::SmallVector<Type, 4> input_types;
  input_types.reserve(num_args);
  // Edit the block arguments and create the new input ops in place to replace
  // the old input ops and quantize ops.

  std::vector<xla::Shape> shapes;
  shapes.reserve(func_type.getNumResults());

  llvm::SmallDenseMap<unsigned, unsigned> output_to_input_alias;
  for (unsigned i = 0; i < main_func.getNumArguments(); ++i)
    if (auto aliasing_output = main_func.getArgAttrOfType<mlir::IntegerAttr>(
            i, "tf.aliasing_output"))
      output_to_input_alias[aliasing_output.getInt()] = i;

      func.front().getTerminator()->eraseOperand(result_idx);
  }

  // Check which function arguments are unused in all branches. We can drop
  // those as well.
  int num_args = cloned_branches[0].getNumArguments();
  llvm::BitVector used_args(num_args);
  for (func::FuncOp func : branches) {
    for (BlockArgument arg : func.getArguments()) {
      if (!arg.use_empty()) used_args.set(arg.getArgNumber());
    }
  }

    SmallPtrSet<Operation*, 4> ops_connected_to_fetch, int num_old_outputs) {
  assert(num_old_outputs + resource_equivalence_classes.getNumClasses() ==
         func.getNumArguments());
  auto graph_op = cast<GraphOp>(func.front().front());

  auto fetch = graph_op.GetFetch();
  OpBuilder builder_chain_src(fetch);
  builder_chain_src.setInsertionPointToStart(fetch->getBlock());

        *module_op->getBodyRegion().getOps().begin());
    if (!func_op) {
      return absl::InternalError("Input MLIR must have only 1 func");
    }
    if (arguments.size() != func_op.getNumArguments()) {
      return absl::InternalError("Input argument has wrong size");
    }

    // Convert input xla::Literal arguments to constants, this allows using
    // constant folding to evaluate function return value.