// preserved and the sequencing rewrite will trigger.
  pm.addPass(mlir::TFDevice::CreateEmbeddingPipeliningPass());
  pm.addPass(mlir::TFDevice::CreateEmbeddingSequencingPass());
  pm.addPass(tensorflow::tf2xla::internal::CreateTPUClusterFormationPass(
      strict_clusters));
  // CreateEmbeddingPipeliningPass may have created more functions, but
  // TPUClusterCleanup and OutsideCompiledToHostLaunch need every function to be

  orig_while_op.erase();

  LOG(INFO) << "EmbeddingPipeliningPass::runOnOperation done.";
}
}  // namespace

std::unique_ptr<OperationPass<mlir::ModuleOp>> CreateEmbeddingPipeliningPass() {
  return std::make_unique<EmbeddingPipeliningPass>();
}

}  // namespace TFDevice