use jax odeint

ajouellette · May 1, 2024 · 5c3a0f7 · 5c3a0f7
1 parent c8399e0
commit 5c3a0f7
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 24 deletions.
diff --git a/jax_cosmo/background.py b/jax_cosmo/background.py
@@ -1,10 +1,10 @@
 # This module implements various functions for the background COSMOLOGY
 import jax.numpy as np
 from jax import lax
+from jax.experimental.ode import odeint
 
 import jax_cosmo.constants as const
 from jax_cosmo.scipy.interpolate import interp
-from jax_cosmo.scipy.ode import odeint
 
 __all__ = [
     "w",

diff --git a/jax_cosmo/core.py b/jax_cosmo/core.py
@@ -1,5 +1,4 @@
 import jax.numpy as np
-from jax.experimental.ode import odeint
 from jax.tree_util import register_pytree_node_class
 
 import jax_cosmo.constants as const

diff --git a/jax_cosmo/scipy/ode.py b/jax_cosmo/scipy/ode.py