* fix test

Joshuaalbert · Aug 14, 2024 · 0c51ac8 · 0c51ac8
1 parent d4948f6
commit 0c51ac8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/essm_jax/tests/test_essm.py b/essm_jax/tests/test_essm.py
@@ -3,9 +3,9 @@
 import jax
 import pytest
 
+jax.config.update('jax_enable_x64', True)
 from essm_jax.sparse import create_sparse_rep, matvec_sparse
 
-jax.config.update('jax_enable_x64', True)
 import numpy as np
 import tensorflow_probability.substrates.jax as tfp
 from jax import numpy as jnp

diff --git a/essm_jax/tests/test_jvp_op.py b/essm_jax/tests/test_jvp_op.py
@@ -1,4 +1,7 @@
+import jax
 import jax.numpy as jnp
+
+jax.config.update('jax_enable_x64', True)
 import numpy as np
 import pytest