Add JAX configuration for disabling async dispatch

brainpy · Oct 25, 2024 · 1792699 · 1792699
1 parent b8e04d8
commit 1792699
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/brainpy/__init__.py b/brainpy/__init__.py
@@ -153,3 +153,8 @@
 
 del deprecation_getattr2
 
+# jax config
+import os
+os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false'
+import jax
+jax.config.update('jax_cpu_enable_async_dispatch', False)