Update dependencies manually

.github/workflows/jax_tests.yml

@@ -13,12 +13,19 @@ jobs:
       fail-fast: false
       matrix:
         jax-version:
-          [0.5.0, 0.5.3, 0.6.1, 0.6.2, 0.7.0, 0.7.2]
-          # 0.4.x versions are not tested because they fail with jax-finufft
-          # 0.5.1 and 0.5.2 installations are broken, see jax#26781
-          # 0.6.0 has a bug with jax.grad, see jax#28144
+          [
+            0.6.2,
+            0.7.0,
+            0.7.2,
+            0.8.0,
+            0.8.1,
+            0.8.3,
+            0.9.1,
+            0.9.2,
+          ]
           # 0.7.0 have performance issues but we still support it, see diffrax#680
           # 0.7.1 fails with equinox, see equinox#1081
+          # 0.8.2 and 0.9.0 are skipped due to interpax_fft
         group: [1, 2]
     steps:
       - uses: actions/checkout@v6

desc/coils.py

@@ -1555,7 +1555,7 @@ def current(self):
     def current(self, new):
         # new must be a 1D iterable regardless of the tree structure of the CoilSet
         old, tree = tree_flatten(self.current)
-        new = jnp.atleast_1d(new).flatten()
+        new = jnp.atleast_1d(jnp.asarray(new)).flatten()
         new = jnp.broadcast_to(new, (len(old),))
         new = tree_unflatten(tree, new)
         for coil, cur in zip(self.coils, new):

desc/particles.py

@@ -15,7 +15,7 @@
 )
 from scipy.constants import Boltzmann, elementary_charge, proton_mass
 
-from desc.backend import jax, jit, jnp, tree_map
+from desc.backend import jax, jnp, tree_map
 from desc.batching import vmap_chunked
 from desc.compute.utils import _compute as compute_fun
 from desc.compute.utils import get_profiles, get_transforms
@@ -130,7 +130,6 @@
         """Coordinate frame of the model."""
         return self._frame
 
-    @jit
     def vf(self, t, x, args):
         """RHS of guiding center trajectories without collisions or slowing down.
 
@@ -415,7 +414,8 @@
         q=2,
     ):
         rho0, theta0, zeta0, xi0, E, m, q = map(
-            jnp.atleast_1d, (rho0, theta0, zeta0, xi0, E, m, q)
+            lambda x: jnp.atleast_1d(jnp.asarray(x)),
+            (rho0, theta0, zeta0, xi0, E, m, q),
         )
         rho0, theta0, zeta0, xi0, E, m, q = jnp.broadcast_arrays(
             rho0, theta0, zeta0, xi0, E, m, q
@@ -547,7 +547,9 @@
         m=4,
         q=2,
     ):
-        R0, phi0, Z0, xi0, E, m, q = map(jnp.atleast_1d, (R0, phi0, Z0, xi0, E, m, q))
+        R0, phi0, Z0, xi0, E, m, q = map(
+            lambda x: jnp.atleast_1d(jnp.asarray(x)), (R0, phi0, Z0, xi0, E, m, q)
+        )
         R0, phi0, Z0, xi0, E, m, q = jnp.broadcast_arrays(R0, phi0, Z0, xi0, E, m, q)
         self.m = m * proton_mass
         self.q = q * elementary_charge
@@ -690,7 +692,7 @@
         is_curve_magnetic_axis=False,
     ):
         self.curve = curve
-        E, m, q = map(jnp.atleast_1d, (E, m, q))
+        E, m, q = map(lambda x: jnp.atleast_1d(jnp.asarray(x)), (E, m, q))
         self.E = jnp.broadcast_to(E, (N,)) * JOULE_PER_EV
         self.m = jnp.broadcast_to(m, (N,)) * proton_mass
         self.q = jnp.broadcast_to(q, (N,)) * elementary_charge
@@ -849,7 +851,7 @@
         is_surface_from_eq=False,
     ):
         self.surface = surface
-        E, m, q = map(jnp.atleast_1d, (E, m, q))
+        E, m, q = map(lambda x: jnp.atleast_1d(jnp.asarray(x)), (E, m, q))
         self.E = jnp.broadcast_to(E, (N,)) * JOULE_PER_EV
         self.m = jnp.broadcast_to(m, (N,)) * proton_mass
         self.q = jnp.broadcast_to(q, (N,)) * elementary_charge

desc/profiles.py

@@ -970,10 +970,10 @@ def __init__(self, values=None, knots=None, method="cubic2", name=""):
 
         if values is None:
             values = [0, 0, 0]
-        values = jnp.atleast_1d(values)
+        values = jnp.atleast_1d(jnp.asarray(values))
         if knots is None:
             knots = jnp.linspace(0, 1, values.size)
-        knots = jnp.atleast_1d(knots)
+        knots = jnp.atleast_1d(jnp.asarray(knots))
         errorif(values.shape[-1] != knots.shape[-1])
         errorif(not (values.ndim == knots.ndim == 1), NotImplementedError)
         self._knots = knots
@@ -1062,10 +1062,10 @@ class HermiteSplineProfile(_Profile):
     def __init__(self, f, df, knots=None, name=""):
         super().__init__(name)
 
-        f, df = jnp.atleast_1d(f, df)
+        f, df = jnp.atleast_1d(jnp.asarray(f), jnp.asarray(df))
         if knots is None:
             knots = jnp.linspace(0, 1, f.size)
-        knots = jnp.atleast_1d(knots)
+        knots = jnp.atleast_1d(jnp.asarray(knots))
         errorif(not (f.shape[-1] == df.shape[-1] == knots.shape[-1]))
         errorif(not (f.ndim == df.ndim == knots.ndim == 1), NotImplementedError)
         self._knots = knots

requirements.txt

@@ -1,7 +1,7 @@
-jax >= 0.5.0, != 0.5.1, != 0.5.2, != 0.6.0, != 0.7.1, < 0.10.0
+jax >= 0.6.2, != 0.7.1, < 0.10.0
 colorama <= 0.4.6
 diffrax >= 0.6.0, <= 0.7.2
-equinox >=0.11.10, <=0.13.7
+equinox >=0.11.10, <=0.13.8
 h5py >= 3.0.0, <= 3.16.0
 interpax >= 0.3.3, < 0.4
 interpax_fft >= 0.0.6, <= 0.0.7
@@ -11,7 +11,7 @@ mpmath >= 1.0.0, <= 1.4.1
 netcdf4 >= 1.5.4, !=1.7.4, <= 1.7.5
 numpy >= 1.20.0, <= 2.5
 nvidia-ml-py >=12.535.77
-optax < 0.3.0
+optax < 0.3
 orthax < 0.3
 plotly >= 5.16, <= 6.7.0
 psutil <= 7.2.2

tests/test_examples.py

@@ -163,6 +163,7 @@ def test_HELIOTRON_vac_results(HELIOTRON_vac):
 @pytest.mark.solve
 def test_solve_bounds():
     """Tests optimizing with bounds=(lower bound, upper bound)."""
+    # Note: This test is known to be sensitive to minor numerical changes
     # decrease resolution and double pressure so no longer in force balance
     eq = get("DSHAPE")
     with pytest.warns(UserWarning, match="Reducing radial"):