void AddCallModuleSerializationPasses(OpPassManager& pm) {
  AddShapeLegalizationPasses(pm);
  pm.addPass(createReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass());
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass may create
  // duplicate constants. Add canonicalizer to deduplicate.
  pm.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
  pm.addPass(TF::CreateXlaCallModuleSerializationPass());
}

  }];
  let dependentDialects = [
      "mlir::func::FuncDialect",
      "mlir::stablehlo::StablehloDialect",
      "TF::TensorFlowDialect",
  ];
}

def ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass : Pass<"stablehlo-replace-stablehlo-ops-in-main-function-with-xla-call-module-ops", "mlir::ModuleOp"> {
  let summary = "Replaces the StableHLO ops with a separate XlaCallModuleOps.";
  let description = [{