void UpdateTerminatorArguments(T& func,
                               const ArrayRef<unsigned> arguments_to_erase,
                               llvm::BitVector& erase_indices) {
  auto terminator = func.front().getTerminator();
  int num_operands = terminator->getNumOperands();
  erase_indices.resize(num_operands);
  for (auto arg_index : arguments_to_erase) {
    auto argument = func.getArgument(arg_index);

  Block& block = region.front();
  if (std::distance(block.begin(), block.end()) != 2) return std::nullopt;
  TF::YieldOp yield =
      llvm::dyn_cast_or_null<TF::YieldOp>(block.getTerminator());
  if (!yield) return std::nullopt;
  func::CallOp call = llvm::dyn_cast_or_null<func::CallOp>(*block.begin());
  if (!call) return std::nullopt;
  if (block.getNumArguments() != call.getNumOperands() ||

  OpBuilder builder(replicate);
  auto new_replicate = builder.create<tf_device::ReplicateOp>(
      replicate.getLoc(), num_replicas, devices, new_replicated_inputs,
      new_packed_inputs,
      replicate.GetBody().getTerminator()->getOperandTypes());
  for (auto arg : replicate.GetBody().getArguments()) {
    if (replicate.IsReplicatedBlockArgument(arg)) {
      arg.replaceAllUsesWith(

    //   ^main(%x : bar):
    //     yield %x : bar
    // }
    //
    // to
    //
    // func.func @f(%x : bar) -> bar {
    //   ...
    // }
    auto ret = f.getBody().front().getTerminator();
    std::vector<Type> argument_types;
    for (auto arg : f.getBody().front().getArguments()) {
      argument_types.push_back(arg.getType());
    }
    std::vector<Type> return_types;

  if (!llvm::hasSingleElement(block.without_terminator())) return false;

  auto graph = llvm::dyn_cast<tf_executor::GraphOp>(block.front());
  if (!graph) return false;

  Operation* terminator = block.getTerminator();
  if (graph.getNumResults() != terminator->getNumOperands()) return false;
  for (auto result : llvm::zip(graph.getResults(), terminator->getOperands()))
    if (std::get<0>(result) != std::get<1>(result)) return false;

  bool IsMaxFunction(Block& block) const {
    if (block.getNumArguments() != 2) return false;

    auto return_op = cast<mlir::stablehlo::ReturnOp>(block.getTerminator());
    if (return_op.getNumOperands() != 1) return false;

    auto max_op = dyn_cast_or_null<MaxOp>(
        return_op.getOperands().front().getDefiningOp());
    if (!max_op) return false;

      if (!passthrough_arg) break;
      value = call.getArgOperands()[passthrough_arg.value()];
    } else if (isa<tf_device::LaunchOp, tf_device::ClusterOp>(op)) {
      value = op->getRegion(0).front().getTerminator()->getOperand(res_index);
    } else {
      break;
    }
  }
  return value;
}

// Analyze the region.
BacktrackAnalysisInfo::BacktrackAnalysisInfo(

    mlir::Operation *terminator = nullptr;
    if (subgraph->name()) {
      if (auto fn = module.lookupSymbol<FuncOp>(subgraph->name()->str()))
        terminator = fn.back().getTerminator();
    }
    i = 0;
    for (auto output : *subgraph->outputs()) {
      print_buffer(*subgraph, i, output, [&](int i) {
        return terminator ? terminator->getOperand(i).getLoc() : unknown_loc;
      });

    arg.replaceAllUsesWith(new_arg);
    arg.dropAllUses();
    bb.eraseArgument(0);
  }

  // Edit the return ops and remove the dequantize ops in place.
  auto* terminator = bb.getTerminator();
  int num_return_operands = terminator->getNumOperands();
  llvm::SmallVector<Type, 4> output_types;
  output_types.reserve(num_return_operands);
  for (int i = 0; i != num_return_operands; ++i) {

    func::FuncOp function,
    llvm::ArrayRef<std::string> var_handle_shared_names) {
  Block& block = function.front();

  auto return_op =
      llvm::dyn_cast_or_null<func::ReturnOp>(block.getTerminator());
  if (!return_op)
    return function.emitError() << "expects function '" << function.getName()
                                << "' to have a MLIR ReturnOp";