fixing 64-bit option in sqgturb

StevePny · May 25, 2024 · 545db67 · 545db67
1 parent a1c3c21
commit 545db67
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dabench/data/sqgturb.py b/dabench/data/sqgturb.py
@@ -36,15 +36,17 @@
 import jax
 import jax.numpy as jnp
 from jax.numpy.fft import rfft2, irfft2
-from jax import config
 from functools import partial
 from importlib import resources
 
 from dabench.data import _data
 from dabench import _suppl_data
 
 # Set to enable 64bit floats in Jax
-config.update('jax_enable_x64', True)
+# Following:
+# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
+import os
+os.environ["JAX_ENABLE_X64"] = 'True'
 
 
 class SQGTurb(_data.Data):