cleanup old jaxlib version check

abaelhe · Jun 7, 2023 · c3d3c19 · c3d3c19
1 parent fa099fd
commit c3d3c19
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
@@ -116,13 +116,7 @@ def _xla_gc_callback(*args):
 xla_extension_version: int = getattr(xla_client, '_version', 0)
 
 import jaxlib.gpu_rnn as gpu_rnn  # pytype: disable=import-error
-if jaxlib.version.__version_info__ >= (0, 4, 11):
-  # TODO(sharadmv): make this unconditional when minimum jaxlib version is
-  # bumped to 0.4.11
-  try:
-    import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error
-  except ModuleNotFoundError:
-    pass
+import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error
 
 # Version number for MLIR:Python APIs, provided by jaxlib.
 mlir_api_version = xla_client.mlir_api_version