Remove leading underscores as in jax-ml/jax@760deb3

eelregit · Mar 31, 2023 · c2f0c24 · c2f0c24
1 parent 346a314
commit c2f0c24
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pmwd/ode_util.py b/pmwd/ode_util.py
@@ -34,7 +34,7 @@
 from jax._src import core
 from jax import custom_derivatives
 from jax import lax
-from jax._src.numpy.util import _promote_dtypes_inexact
+from jax._src.numpy.util import promote_dtypes_inexact
 from jax._src.util import safe_map, safe_zip
 from jax.flatten_util import ravel_pytree
 from jax.tree_util import tree_leaves, tree_map
@@ -80,7 +80,7 @@ def initial_step_size(fun, t0, y0, order, rtol, atol, f0):
     # Algorithm from:
     # E. Hairer, S. P. Norsett G. Wanner,
     # Solving Ordinary Differential Equations I: Nonstiff Problems, Sec. II.4.
-    y0, f0 = _promote_dtypes_inexact(y0, f0)
+    y0, f0 = promote_dtypes_inexact(y0, f0)
     dtype = y0.dtype
 
     scale = atol + jnp.abs(y0) * rtol

diff --git a/setup.py b/setup.py
@@ -24,7 +24,7 @@
     packages=find_packages(),
     python_requires='>=3.7',
     install_requires=[
-        'jax',
+        'jax>=0.4.7',
         'numpy>=1.20',  # numpy.typing
         'mcfit',
     ],