}

LogicalResult ModifyIONodesPass::ModifyOutputNodes(
    func::FuncOp func, llvm::SmallVectorImpl<Type>& new_output_types,
    OpBuilder builder) {
  Block& block = func.front();
  auto* terminator = block.getTerminator();
  builder.setInsertionPoint(terminator);

  if (mlir::isa<FloatType>(output_type)) {
    return success();
  }

  int num_return_operands = terminator->getNumOperands();

    //   ^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(

LogicalResult IfRegionOp::verifyRegions() {
  IfRegionOp op = *this;
  TypeRange then_types =
      op.getThenBranch().front().getTerminator()->getOperandTypes();
  TypeRange else_types =
      op.getElseBranch().front().getTerminator()->getOperandTypes();

  TypeRangeWithDesc results{op.getResultTypes(), "result"};
  TypeRangeWithDesc then_results{then_types, "then result"};

  }
  return changed;
}

bool ShapeInference::InferShapeForIfRegion(IfRegionOp op) {
  bool changed = false;

  Operation* then_yield = op.getThenBranch().front().getTerminator();
  Operation* else_yield = op.getElseBranch().front().getTerminator();
  for (auto result : zip(op.getResults(), then_yield->getOperandTypes(),
                         else_yield->getOperandTypes())) {

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