SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
    llvm::StringMap<Attribute>* derived_attrs) const {
  for (const auto& operand :
       llvm::enumerate(signature.getFunctionType().getInputs())) {
    // If the index is larger than the operand number of the call_op, the
    // default value of the operand needs to be used.
    if (operand.index() >= call_op.getNumOperands()) {

      rewriter.create<func::FuncOp>(input_op->getLoc(), func_name, func_type);

  OpBuilder::InsertionGuard guard = OpBuilder::InsertionGuard(rewriter);
  ArrayRef<Type> inputs = quantization_func.getFunctionType().getInputs();
  Block* block = rewriter.createBlock(
      &quantization_func.getBody(), quantization_func.begin(), inputs,
      SmallVector<Location>(inputs.size(), quantization_func.getLoc()));

    mlir::Region *getCallableRegion();

    /// Returns the argument types of this function.
    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }

    /// Returns the result types of this function.
    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }

    // Get the names of all defined attributes, including both derived and
    // non-derived ones.

    for (auto ret_op : ret->getOperands()) {
      return_types.push_back(ret_op.getType());
    }
    auto newType =
        FunctionType::get(rewriter.getContext(), argument_types, return_types);
    if (f.getFunctionType() == newType) {
      return failure();
    }
    rewriter.modifyOpInPlace(f, [&] { f.setType(newType); });

    // Adjust the type of the return values callers of the function to

  // at the end?
  int first_tensor = -1, last_tensor = -1, first_tensor_list = -1,
      last_tensor_list = -1, first_attr = -1;
  for (auto arg : llvm::enumerate(func.getFunctionType().getInputs())) {
    Type arg_type = arg.value();

    if (auto tensor = arg_type.dyn_cast<TFRTensorType>()) {
      if (first_tensor == -1) {
        first_tensor = arg.index();
      }

    func.getRegion().addArgument(chaining_data_type, func.getLoc());
  }

  FunctionType ftype =
      FunctionType::get(func.getContext(), func.getBody().getArgumentTypes(),
                        func.getFunctionType().getResults());
  func.setType(ftype);
}

// Appends function results with `num_resources` number of results of requested
// type.
void AppendFunctionResults(func::FuncOp func, int num_resources,

  SmallVector<Type> input_types(main_func_op.getArgumentTypes());
  input_types.emplace_back(file_prefix_arg_type);

  main_func_op.setType(
      builder.getFunctionType(input_types, main_func_op.getResultTypes()));

  // Add "__tf_file_prefix" to the "tf_saved_model.index_path" attribute for the
  // newly created argument.
  main_func_op.setArgAttr(new_file_prefix_arg.getArgNumber(),

        (*stats)[grad_var].accumulate_on_write = true;
      }
    }
  }
  OpBuilder builder(while_op);
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getInputs(), operands,
      while_op->getAttrs());
  for (int64_t i = 0; i < while_op.getNumOperands(); ++i) {
    if (ta_arg_buffer_type(i)) {
      while_op.getResult(i).replaceAllUsesWith(while_op.getOperand(i));

  mlir::func::FuncOp main_func =
      module.lookupSymbol<mlir::func::FuncOp>("main");
  TF_RET_CHECK(main_func != nullptr) << "No main function found";
  mlir::FunctionType func_type = main_func.getFunctionType();

  int num_args = func_type.getNumInputs();
  xla_input_shapes->reserve(num_args);

  std::vector<xla::Shape> individual_arg_shapes;
  individual_arg_shapes.reserve(num_args);

    // 3. Captures inputs.
    // The number of dataset inputs can be indirectly determined to be
    // total_number_of_inputs - state_inputs - captured_inputs.
    auto func_inputs = reduce_func.getFunctionType().getInputs();
    const int func_input_size = func_inputs.size();
    const int argument_size =
        reduce_dataset->getAttrOfType<ArrayAttr>("Targuments").size();
    const int state_size =