Fix jax config (#301)

ratt-ru · Apr 4, 2024 · d7d7b6c · d7d7b6c
1 parent 0d60d8e
commit d7d7b6c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -4,6 +4,7 @@ History
 
 X.Y.Z (YYYY-MM-DD)
 ------------------
+* Fix jax configuration (:pr:`301`)
 * Configure dependabot for github actions (:pr:`296`)
 * Cache numba kernels between CI runs (:pr:`294`)
 

diff --git a/africanus/util/jax_init.py b/africanus/util/jax_init.py
@@ -7,7 +7,7 @@
 
 
 try:
-    import jax.config
+    import jax
 except ImportError:
     pass
 else: