suppress warnings for jax initialize backend

sokrypton · Jan 23, 2024 · 164f3f8 · 164f3f8
1 parent 941feec
commit 164f3f8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/colabfold/batch.py b/colabfold/batch.py
@@ -75,7 +75,11 @@
 logger = logging.getLogger(__name__)
 import jax
 import jax.numpy as jnp
-logging.getLogger('jax._src.lib.xla_bridge').addFilter(lambda _: False)
+
+# from jax 0.4.6, jax._src.lib.xla_bridge moved to jax._src.xla_bridge
+# suppress warnings: Unable to initialize backend 'rocm' or 'tpu'
+logging.getLogger('jax._src.xla_bridge').addFilter(lambda _: False) # before jax 0.4.5
+logging.getLogger('jax._src.lib.xla_bridge').addFilter(lambda _: False) # from jax 0.4.6
 
 def mk_mock_template(
     query_sequence: Union[List[str], str], num_temp: int = 1