// Propagate the analysis results from within the callee's body.
      PropagatePotentiallyWrittenUpFromCallee(batch_function.func().getRegion(),
                                              batch_function.getOperands());
      return;
    }
    // For all other ops, we assume it mutates all resources it uses, so
    // this errs on the side of being conservative. We should improve

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

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

    Operation::operand_range mean(fused_batch_norm->getOperands());
    ::mlir::FloatAttr exponential_avg_factor;
    ::mlir::TF::FusedBatchNormV3Op root;
    Operation::operand_range offset(fused_batch_norm->getOperands());
    Operation::operand_range x(fused_batch_norm->getOperands());
    Operation::operand_range scale(fused_batch_norm->getOperands());
    Operation::operand_range variance(fused_batch_norm->getOperands());

      island.getBody().front().begin(), body.getOperations(),
      copy_range.begin(), copy_range.end());
  builder.setInsertionPointToEnd(&island.getBody().front());
  builder.create<tf_executor::YieldOp>(loc, return_op.getOperands());
  for (auto item : llvm::enumerate(graph_op.getResults())) {
    return_op.setOperand(item.index(), item.value());
  }
}

std::unique_ptr<OperationPass<func::FuncOp>>

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

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