Disable async dispatch on JAX CPU for a jax_cfd test, i.e., opt out o…

…f the change introduced in jax-ml/jax#20470, as it is causing some real failures under tsan. PiperOrigin-RevId: 625707302
google · Apr 17, 2024 · 214b831 · 214b831
1 parent d215f13
commit 214b831
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/jax_cfd/ml/equations_test.py b/jax_cfd/ml/equations_test.py
@@ -293,4 +293,7 @@ def test_epd_modules(
 
 
 if __name__ == '__main__':
+  # Temporarily disable async dispatch on JAX CPU due to tsan error.
+  jax.config.update('jax_cpu_enable_async_dispatch', False)
+
   absltest.main()