// 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();
      }

      // The FuncOp type can contain types that the op's operand and result
      // types do not contain.
      if (auto func = dyn_cast<func::FuncOp>(op)) {
        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
      } else if (auto bitcast_convert_op =
                     dyn_cast<mlir::stablehlo::BitcastConvertOp>(op)) {
        return IsConvertOpLegal<mlir::stablehlo::BitcastConvertOp,

    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,

      TypesFromValues(subgraph.FuncArguments());
  llvm::SmallVector<Type> return_types =
      TypesFromValues(subgraph.FuncOutputs());

  FunctionType function_type =
      builder.getFunctionType(input_types, return_types);

  std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);

  func::FuncOp new_func = func::FuncOp::create(builder.getUnknownLoc(),

  auto filename_input_type = RankedTensorType::get(
      /*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());

  FunctionType func_type = builder.getFunctionType(
      /*inputs=*/{filename_input_type}, /*results=*/{});
  auto save_func = builder.create<func::FuncOp>(
      NameLoc::get(builder.getStringAttr(kTfQuantSaveFuncName)),

  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));

                            ConversionPatternRewriter *rewriter) const {
    auto guard = OpBuilder::InsertionGuard(*rewriter);
    auto inputs = branch_func.getFunctionType().getInputs();
    Block *block = rewriter->createBlock(
        &branch_func.getBody(), branch_func.begin(), inputs,
        SmallVector<Location>(inputs.size(), branch_func.getLoc()));

  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 =