// TODO(b/284222309): Support lifting for functions called by control flow.
bool UsedBySupportedOps(ModuleOp module, func::FuncOp func) {
  auto function_uses =
      SymbolTable::getSymbolUses(func, &module.getBodyRegion());
  if (!function_uses.has_value()) return false;
  for (auto& function_use : function_uses.value()) {
    if (!llvm::isa<TF::PartitionedCallOp, TF::StatefulPartitionedCallOp>(