auto arg_types = func.front().getArgumentTypes();
  auto result_types = func.front().getTerminator()->getOperandTypes();
  func.setType(FunctionType::get(func.getContext(), arg_types, result_types));
}

void HandleFuncOp(Operation* op) {
  auto func = llvm::cast<func::FuncOp>(op);
  UpdateFuncType(func);
}

// Handles cast op between the first convolution and the block argument.