[mosaic_gpu] Removed debug prints in jax._src.lib.mosaic_gpu

PiperOrigin-RevId: 717913893
jax-ml · Jan 21, 2025 · dd773b9 · dd773b9
1 parent c4643c6
commit dd773b9
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/jax/_src/lib/mosaic_gpu.py b/jax/_src/lib/mosaic_gpu.py
@@ -17,12 +17,9 @@
 try:
   try:
     from jaxlib.mosaic.gpu import _mosaic_gpu_ext  # pytype: disable=import-error
-  except ImportError as e:
-    print(e)
+  except ImportError:
     from jax_cuda12_plugin import _mosaic_gpu_ext  # pytype: disable=import-error
-except ImportError as e:
-  print("="*128)
-  print(e)
+except ImportError:
   raise ModuleNotFoundError("Failed to import the Mosaic GPU bindings") from e
 else:
   _mosaic_gpu_ext.register_passes()