// Don't prune if result is used.
  if (!result.use_empty()) return false;

  Block& body_block = while_op.getBody().front();
  Block& cond_block = while_op.getCond().front();
  Operation* body_yield_op = body_block.getTerminator();

  // The body yield operand, body block argument, condition block argument, and
  // result all correspond to each other (see definition of `WhileRegionOp`).
  int idx = result.getResultNumber();

    if (!func_op) return failure();
    // The quantized fusion should have requantize and return ops at the end.
    auto return_op = dyn_cast_or_null<func::ReturnOp>(
        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();
    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
        return_op.getOperands()[0].getDefiningOp());
    if (!req_op) return failure();

  const int size_before = output_types->size();
  for (int index = start; index < end; ++index) {
    Block& block = parallel_execute.GetRegionBlockWithIndex(index);
    auto terminator_operand_types = block.getTerminator()->getOperandTypes();
    output_types->append(terminator_operand_types.begin(),
                         terminator_operand_types.end());
  }
  return output_types->size() - size_before;
}

                                            tf_device::LaunchOp launch) {
  // Forward launch inner op results to launch op results.
  launch.replaceAllUsesWith(launch.GetBody().getTerminator()->getOperands());

  // For all inner ops, assign the launch device as a `device` attribute.
  if (failed(AssignDevicesInRegion(tf_dialect, launch, launch.getBody())))
    return failure();

    old_val.replaceAllUsesWith(new_val);
    return;
  }
  Operation* old_op = old_val.getDefiningOp();
  Operation* terminator_op =
      old_op->getParentOfType<func::FuncOp>().front().getTerminator();
  llvm::SmallPtrSet<mlir::Operation*, 1> exceptions = {terminator_op};
  old_val.replaceAllUsesExcept(new_val, exceptions);
}

struct TensorArrayStats {

// Returns the aliasing argument number of a fucntion return value if it simply
// forwards the argument. Otherwise, returns -1.
int64_t FindAliasedInput(func::FuncOp func, int64_t return_index) {
  Value return_val = func.front().getTerminator()->getOperand(return_index);
  auto maybe_arg = mlir::dyn_cast<BlockArgument>(return_val);
  if (!maybe_arg) return -1;
  return maybe_arg.getArgNumber();
}

// op operands' types.
//
// Requires the function has exactly one block.
void UpdateFuncType(func::FuncOp func) {
  Operation* terminator = func.front().getTerminator();
  auto return_types = llvm::to_vector<4>(terminator->getOperandTypes());

  FunctionType func_type = func.getFunctionType();
  if (llvm::ArrayRef(return_types) == func_type.getResults()) return;

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

  // Replace `tf_device.return` terminator with `std.return` in function
  // body.
  auto return_op =
      cast<tf_device::ReturnOp>(outlined_func_block->getTerminator());
  builder->setInsertionPoint(return_op);
  ReplaceClusterReturnWithReturn(return_op, builder);

  symbol_table->insert(outlined_func);
  return outlined_func;
}

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