// new passes generally should go into the internal
  // AddReplicatedBridgeClusteringPipelinePasses.
  pm.addPass(mlir::TFTPU::CreateTPUValidateInputsPass());
  pm.addNestedPass<FuncOp>(
      mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
  tensorflow::tf2xla::internal::AddReplicatedBridgeClusteringPipelinePasses(
      pm, module_name);
}

namespace tensorflow {
namespace tf2xla {
namespace internal {

using mlir::OpPassManager;

TEST(ClusteringBridgePassesTest, AddsBridgePasses) {
  OpPassManager pass_manager;
  AddReplicatedBridgeClusteringPipelinePasses(pass_manager);

  EXPECT_EQ(pass_manager.size(), 45);
}

TEST(ClusteringBridgePassesTest, AddsNonTPUBridgePasses) {
  OpPassManager pass_manager;

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Given the pass manager, add Bridge passes to cluster the replicated input
// graphs.
void AddReplicatedBridgeClusteringPipelinePasses(
    mlir::OpPassManager& pm, llvm::StringRef module_name = llvm::StringRef());

// Same as above but for non replicated graphs.
void AddNonReplicatedBridgeClusteringPipelinePasses(mlir::OpPassManager& pm);

    num_submodules++;
    clustering_pipeline_status = RunTFXLABridge(
        submodule,
        [](OpPassManager &pm) {
          internal::AddReplicatedBridgeClusteringPipelinePasses(pm);
        },
        /*module_name=*/"", /*dump_prefix=*/"tf_xla_clustering_bridge_v1");
    if (num_submodules > 1) {
      return mlir::WalkResult::interrupt();
    }

namespace tf2xla {
namespace internal {

using mlir::OpPassManager;
using mlir::func::FuncOp;

// Adds replicated Bridge clustering pipeline passes to the given pass_manager.
// Does not run them.
void AddReplicatedBridgeClusteringPipelinePasses(OpPassManager& pm,
                                                 llvm::StringRef module_name) {
  // The following ops must be preserved regardless of reachability. Ideally,