*count = output_type.getNumElements();
    return true;
  }

  static bool GetInputTensorTotalSize(mlir::Operation* op, int64_t* count) {
    int64_t total_count = 0;
    for (auto input : op->getOperands()) {
      auto input_type =
          mlir::dyn_cast_or_null<mlir::RankedTensorType>(input.getType());
      if (!input_type || !input_type.hasStaticShape()) {
        return false;
      }

  auto old_terminator = block.getTerminator();
  auto new_outputs = llvm::to_vector<8>(old_terminator->getOperands());
  llvm::SmallVector<std::tuple<int64_t, int64_t, bool>, 8>
      output_buffer_to_size;
  for (auto retval : llvm::enumerate(old_terminator->getOperands())) {
    auto it = buffer_to_size.find(retval.value());
    if (it == buffer_to_size.end()) continue;

        composite->getLoc(), composite->getResultTypes(),
        composite->getOperands().slice(2, 3), composite.getName(),
        CustomOption(&builder, custom_option_buffer));
  }
  return builder.create<TFL::CustomOp>(
      composite->getLoc(), composite->getResultTypes(),
      composite->getOperands(), composite.getName(),
      CustomOption(&builder, custom_option_buffer));
}

}  // namespace

// Get total bytes transferred.
int64_t GetTransferredTensorBytes(func::CallOp from_graph,
                                  func::CallOp to_graph) {
  int64_t total_size_transferred = 0;
  for (auto input : to_graph.getOperands()) {
    Operation* input_op = input.getDefiningOp();
    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());

  YieldOp yield = GetYield();
  return wrapped_op.getNumResults() == yield.getNumOperands() &&
         std::equal(wrapped_op.getResults().begin(),
                    wrapped_op.getResults().end(), yield.getOperands().begin());
}

mlir::LogicalResult IslandOp::verify() {
  IslandOp island = *this;
  if (!island.GetBody().args_empty())
    return island.emitOpError() << "expects body without any arguments";

        op.getLoc(), op.getType(), op.getOperands(),
        builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),
        builder.getI64ArrayAttr(partition_dims), op.get_XlaShardingAttr());

      liveouts.snapshot_previous_state();
      return;
    }

    reverse_subgraph.insert(op);
    defined_values.insert(op->getResults().begin(), op->getResults().end());
    operands.insert(op->getOperands().begin(), op->getOperands().end());
  };

  for (Operation* op : reverse_main_func_block_ops) {
    if (!ops_to_add.contains(op)) continue;
    // When hitting a non-StableHLO op, i.e. tf.CustomAggregatorOp, start

    return std::nullopt;
  for (auto [arg, operand] :
       llvm::zip(block.getArguments(), call.getOperands())) {
    if (arg != operand) return std::nullopt;
  }
  for (auto [ret, operand] :
       llvm::zip(call.getResults(), yield.getOperands())) {
    if (ret != operand) return std::nullopt;
  }
  SymbolRefAttr symbol = call.getCallableForCallee().get<SymbolRefAttr>();

  llvm::SmallVector<Operation*> predecessors;
  for (Value operand : op->getOperands()) {
    if (Operation* pred = operand.getDefiningOp()) {
      int result_number = llvm::cast<OpResult>(operand).getResultNumber();
      if (auto pred_island = llvm::dyn_cast<tf_executor::IslandOp>(pred)) {
        Value yield_operand = pred_island.GetYield().getOperand(result_number);
        predecessors.push_back(yield_operand.getDefiningOp());
      }

      rewriter.setInsertionPoint(current_terminator);
      SmallVector<Value, 4> scf_yield_input;
      for (auto it : llvm::zip(tf_if_region_return_type,
                               current_terminator->getOperands())) {
        scf_yield_input.push_back(rewriter.create<CastOp>(
            current_terminator->getLoc(), std::get<0>(it), std::get<1>(it)));
      }

      rewriter.replaceOpWithNewOp<scf::YieldOp>(current_terminator,