public:
  explicit TPUClusterFormationPass(bool strict_clusters)
      : strict_clusters_(strict_clusters) {}

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<mlir::tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;

 private:
  bool strict_clusters_;
};

CreateVerifyClusteringPass();

// Creates a pass that forms clusters from operations of the same
// `_replication_info` attribute.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateTPUClusterFormationPass(bool strict_clusters = false);

// Creates a pass that extracts outside compilation (Host ops inside device
// cluster) at head/tail of Device cluster to run before/after XLA computation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>

  // currently not the case (see b/177478741).
  const llvm::SmallVector<std::string, 4> ops_to_preserve = {
      "tf.TPUReplicateMetadata", "tf.TPUCompilationResult",
      "tf.TPUReplicatedOutput"};
  bool strict_clusters =
      tensorflow::GetMlirCommonFlags()->tf_mlir_enable_strict_clusters;
  pm.addNestedPass<FuncOp>(
      mlir::tf_executor::CreateTFExecutorGraphPruningPass(ops_to_preserve));