}
      }
    }
  }
  return LogicalResult::success();
}

struct Callers {
  TF::StatefulPartitionedCallOp forward;
  TF::StatefulPartitionedCallOp core_tpu;
  TF::StatefulPartitionedCallOp backward;
  TF::StatefulPartitionedCallOp non_tpu;
};

template <typename InputContainer>
std::vector<Type> GetValueTypes(const InputContainer& input) {

// cluster) ops to a separate parallel_execute region to run on CPU.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractOutsideCompilationPass();

// Create a pass that encapsulates StatefulPartitionedCallOp within a cluster.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateXlaClusterFormationPass();

// Creates a pass that marks unsupported ops in device cluster for outside
// compilation.

            operations->insert(successor);
          }
        }
      }
    }
    ops_to_process.swap(new_ops);
    new_ops.clear();
  }
}

TF::StatefulPartitionedCallOp MakeFuncCaller(
    mlir::OpBuilder& builder, const Location& loc, func::FuncOp func,
    const llvm::SetVector<Value>& operands) {
  // Constructs a tf.StatefulPartitionedCall to the function provided in 'func'

  EXPECT_EQ(entry_funcs[0].getSymName(), "entry_func");
  llvm::SmallVector<mlir::SymbolUserOpInterface> outermost_pcall_ops;
  auto result =
      mlir::GetFirstOpsOfType<mlir::TF::StatefulPartitionedCallOp,
                              mlir::TF::PartitionedCallOp>(
          entry_funcs[0], symtab, has_compile_device_type, outermost_pcall_ops)
          .succeeded();
  ASSERT_TRUE(result);

  llvm::SmallVector<SymbolUserOpInterface> noinline_pcall_ops,
      outermost_pcall_ops;
  if (mlir::failed(
          mlir::GetOpsOfTypeUntilMiss<mlir::TF::StatefulPartitionedCallOp,
                                      mlir::TF::PartitionedCallOp>(
              func, symtab, /*predicate*/ has_no_compile_device_type,
              /*hits*/ noinline_pcall_ops,

    target.addLegalOp<TF::ResizeBilinearOp>();
    target.addLegalOp<TF::ResizeNearestNeighborOp>();
  }
  if (skip_partitioned_calls_) {
    target.addLegalOp<TF::PartitionedCallOp>();
    target.addLegalOp<TF::StatefulPartitionedCallOp>();
  }

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  if (failed(applyPartialConversion(func, target, frozen_patterns))) {
    return signalPassFailure();
  }
}

      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>(
            function_use.getUser())) {
      return false;
    }
  }
  return true;
}

// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.

    runOnOperation() {
  ModuleOp module_op = getOperation();

  func::FuncOp main_func = FindMainFuncOp(module_op);
  if (!main_func) return;

  // In case the model has tf.StatefulPartitionedCallOp or tf.PartitionedCallOp,
  // we recursively find called functions and process StableHLO ops in them.
  SmallVector<func::FuncOp> func_ops;
  func_ops.push_back(main_func);
  int stablehlo_func_id = -1;

bool IsAlwaysAllowlistedOp(Operation *op) {
  return llvm::isa<
      // clang-format off
      // go/keep-sorted start
      TF::ConstOp,
      TF::IdentityOp,
      TF::PartitionedCallOp,
      TF::StatefulPartitionedCallOp
      // go/keep-sorted end
      // clang-format on
      >(op);
}

// LINT.IfChange
// The list of quantizable ops in the Legacy Integer mode.
ABSL_ATTRIBUTE_NOINLINE const std::set<std::string>

// subcomputation in the TF/XLA bridge.
bool SupportsCommunicationComputation(Operation* op) {
  return isa<TF::IfRegionOp, TF::WhileRegionOp, TF::CaseRegionOp,
             TF::XlaCallModuleOp, TF::StatefulPartitionedCallOp,
             TF::PartitionedCallOp, TF::LegacyCallOp>(op);
}

#define GEN_PASS_DEF_PREPARETPUCOMPUTATIONFORTFEXPORTPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"