From e15eb9d17370c5daddb8fdff952f44fdf0286607 Mon Sep 17 00:00:00 2001 From: Steven Chiu Date: Sat, 20 Jul 2024 00:55:46 -0500 Subject: [PATCH] ImportError Fixed: Use jax.config.update and remove import .__version (#138) * use jax.config.update * add version import back --- tornadox/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tornadox/__init__.py b/tornadox/__init__.py index 4b0548c..28d5db7 100644 --- a/tornadox/__init__.py +++ b/tornadox/__init__.py @@ -1,10 +1,10 @@ """Collect all modules into the tornadox.* namespace""" -from jax.config import config +import jax from . import ek0, ek1, experimental, init, ivp, iwp, kalman, odefilter, rv, sqrt, step -config.update("jax_enable_x64", True) +jax.config.update("jax_enable_x64", True) -from ._version import version as __version__ +from ._version import version as __version__ \ No newline at end of file