];
}

def LegalizeTFCollective : Pass<"xla-legalize-tf-collective", "ModuleOp"> {
  let summary = "Legalize TF/XLA collective ops (TensorFlow dialect) to the HLO dialect";

  let constructor = "mlir::mhlo::CreateLegalizeTFCollectivePass()";
  let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
}

def VerifyTFXLALegalization : Pass<"tfxla-verify-legalization", "mlir::func::FuncOp"> {

// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();

// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();

// Verifies that the TF/XLA ops have all been lowered to MHLO.
std::unique_ptr<OperationPass<func::FuncOp>> CreateVerifyTFXLALegalizationPass(
    bool legalize_chlo = true);

  if (failed(applyPatternsAndFoldGreedily(module, std::move(patterns)))) {
    signalPassFailure();
  }
}
}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass() {
  return std::make_unique<LegalizeTFCollective>();
}

}  // namespace mhlo

  if (lower_to_xla_hlo) {
    for (auto& target_pass : custom_legalization_passes) {
      pm.addNestedPass<mlir::func::FuncOp>(std::move(target_pass));
    }
    pm.addPass(mlir::mhlo::CreateLegalizeTFCollectivePass());
  }

  // These passes are grouped together as they have to run in specific order.
  // Passes before this can run relatively in any order, as long as they happen
  // before legalization.