diff --git a/jaxlib/mosaic/gpu/custom_call.cc b/jaxlib/mosaic/gpu/custom_call.cc
index 5a3facbb2..361c839b6 100644
--- a/jaxlib/mosaic/gpu/custom_call.cc
+++ b/jaxlib/mosaic/gpu/custom_call.cc
@@ -130,15 +130,15 @@ mlir::FailureOr<mlir::OpPassManager> GetPassPipeline(
     mlir::registerStripDebugInfo();
     mlir::registerConvertNVGPUToNVVMPass();
     mlir::registerConvertVectorToSCF();
-    mlir::registerSCFToControlFlow();
+    mlir::registerSCFToControlFlowPass();
     mlir::registerConvertNVVMToLLVMPass();
     mlir::registerArithToLLVMConversionPass();
     mlir::registerConvertIndexToLLVMPass();
     mlir::registerConvertGpuOpsToNVVMOps();
     mlir::registerConvertMathToLLVMPass();
     mlir::registerConvertFuncToLLVMPass();
-    mlir::registerConvertAffineToStandard();
-    mlir::registerReconcileUnrealizedCasts();
+    mlir::registerLowerAffinePass();
+    mlir::registerReconcileUnrealizedCastsPass();
     // TODO(apaszke): Only register the passes we actually use.
     mlir::memref::registerMemRefPasses();
     mlir::registerConvertToLLVMPass();