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

      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,

// Returns the predecessors of `op` when `op`'s predecessors are wrapped by
// islands.
llvm::SmallVector<Operation*> IslandPredecessors(Operation* op) {
  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)) {

      func_operand_types.push_back(operand.getType());

    // Function results are the yield operands
    SmallVector<Type, 16> func_result_types;
    for (Value operand : island_op.GetYield().getOperands())
      func_result_types.push_back(operand.getType());
    FunctionType func_type =
        FunctionType::get(ctx, func_operand_types, func_result_types);

    // Create the outlined function

  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;

  return true;
}

// Checks if an operation of the tf_executor dialect can have TPU devices
// propagated through.

    partition_dims[partition_dim] = num_cores_per_replica.getInt();
  }

  if constexpr (is_input) {
    auto pi = builder.create<TF::TPUPartitionedInputV2Op>(
        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>(

                 OpBuilder &builder) {
  llvm::SmallVector<Value> non_resource_args, resource_args;
  bool has_resources = false, in_order = true;
  for (const Value &arg : cluster_func_op.getOperands()) {
    if (!mlir::isa<TF::ResourceType>(getElementTypeOrSelf(arg.getType()))) {
      non_resource_args.push_back(arg);
      if (has_resources) in_order = false;
    } else {
      resource_args.push_back(arg);

  module.walk([&](tf_device::LaunchFuncOp launch) {
    OpBuilder builder(launch);
    auto call_op = builder.create<TF::PartitionedCallOp>(
        module.getLoc(), launch.getResultTypes(), launch.getOperands(),
        SymbolRefAttr::get(builder.getContext(), launch.getFunc()),
        /*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));

    Operation* terminator = execute_block.getTerminator();
    builder->setInsertionPoint(terminator);
    auto yield = builder->create<tf_executor::YieldOp>(
        terminator->getLoc(), terminator->getOperands());
    terminator->erase();

    // Create new island for each region.
    builder->setInsertionPoint(island_op);
    auto execute_island = builder->create<tf_executor::IslandOp>(

    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {
          this->join(getLatticeElement(arg), *getLatticeElement(value));
        }
      }
    } else if (auto while_op = dyn_cast<TF::WhileOp>(op)) {