From d190dea92a811d0c32eb9b8d97db75f9f40500a9 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 20 May 2026 01:38:00 -0400 Subject: [PATCH 1/9] update dependencies manually --- .github/workflows/jax_tests.yml | 17 ++++++++++++++++- requirements.txt | 6 +++--- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index 3736ab63b5..58c2747250 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -13,7 +13,22 @@ 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.5.0, + 0.5.3, + 0.6.1, + 0.6.2, + 0.7.0, + 0.7.2, + 0.8.0, + 0.8.1, + 0.8.2, + 0.8.3, + 0.9.0, + 0.9.1, + 0.9.2, + 0.10.0, + ] # 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 diff --git a/requirements.txt b/requirements.txt index 80b282e29d..671f5dd861 100644 --- a/requirements.txt +++ b/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.5.0, != 0.5.1, != 0.5.2, != 0.6.0, != 0.7.1, < 0.11.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 From f312efc617f086edd7e33e480694cb8a144fcb61 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 20 May 2026 12:01:48 -0400 Subject: [PATCH 2/9] fix jnp.atleast_1d errors --- .github/workflows/jax_tests.yml | 2 -- desc/coils.py | 2 +- desc/particles.py | 5 ++++- desc/profiles.py | 8 ++++---- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index 58c2747250..1699eafa64 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -22,9 +22,7 @@ jobs: 0.7.2, 0.8.0, 0.8.1, - 0.8.2, 0.8.3, - 0.9.0, 0.9.1, 0.9.2, 0.10.0, diff --git a/desc/coils.py b/desc/coils.py index 7faa805e53..42f96b1ce8 100644 --- a/desc/coils.py +++ b/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): diff --git a/desc/particles.py b/desc/particles.py index 8a3db05929..8885ea4abb 100644 --- a/desc/particles.py +++ b/desc/particles.py @@ -415,7 +415,7 @@ def __init__( q=2, ): rho0, theta0, zeta0, xi0, E, m, q = map( - jnp.atleast_1d, (rho0, theta0, zeta0, xi0, E, m, q) + jnp.asarray, (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,6 +547,7 @@ def __init__( m=4, q=2, ): + R0, phi0, Z0, xi0, E, m, q = map(jnp.asarray, (R0, phi0, Z0, xi0, E, m, q)) 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 = jnp.broadcast_arrays(R0, phi0, Z0, xi0, E, m, q) self.m = m * proton_mass @@ -690,6 +691,7 @@ def __init__( is_curve_magnetic_axis=False, ): self.curve = curve + E, m, q = map(jnp.asarray, (E, m, q)) E, m, q = map(jnp.atleast_1d, (E, m, q)) self.E = jnp.broadcast_to(E, (N,)) * JOULE_PER_EV self.m = jnp.broadcast_to(m, (N,)) * proton_mass @@ -849,6 +851,7 @@ def __init__( is_surface_from_eq=False, ): self.surface = surface + E, m, q = map(jnp.asarray, (E, m, q)) E, m, q = map(jnp.atleast_1d, (E, m, q)) self.E = jnp.broadcast_to(E, (N,)) * JOULE_PER_EV self.m = jnp.broadcast_to(m, (N,)) * proton_mass diff --git a/desc/profiles.py b/desc/profiles.py index 3fd09149ed..ded35cd32f 100644 --- a/desc/profiles.py +++ b/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 From 56c6fd4f818316bc08a1f47bed0cffefa13e8ad1 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 20 May 2026 12:07:32 -0400 Subject: [PATCH 3/9] update dependency explanations --- .github/workflows/jax_tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index 1699eafa64..dbe1d562a8 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -27,11 +27,11 @@ jobs: 0.9.2, 0.10.0, ] - # 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.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 From 4358dcf488ab9adf0379abc170295b32ac3c31c0 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 20 May 2026 13:59:30 -0400 Subject: [PATCH 4/9] fix --- desc/particles.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/desc/particles.py b/desc/particles.py index 8885ea4abb..852522e96d 100644 --- a/desc/particles.py +++ b/desc/particles.py @@ -417,6 +417,9 @@ def __init__( rho0, theta0, zeta0, xi0, E, m, q = map( jnp.asarray, (rho0, theta0, zeta0, xi0, E, m, q) ) + rho0, theta0, zeta0, xi0, E, m, q = map( + jnp.atleast_1d, (rho0, theta0, zeta0, xi0, E, m, q) + ) rho0, theta0, zeta0, xi0, E, m, q = jnp.broadcast_arrays( rho0, theta0, zeta0, xi0, E, m, q ) From ed2ccd09f9e8cf7620327159f4b8e1b2004422ef Mon Sep 17 00:00:00 2001 From: YigitElma Date: Fri, 29 May 2026 00:44:38 +0300 Subject: [PATCH 5/9] bump min jax version --- .github/workflows/jax_tests.yml | 4 +--- requirements.txt | 2 +- tests/test_examples.py | 1 + 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index dbe1d562a8..e14bdd79e0 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -14,9 +14,6 @@ jobs: matrix: jax-version: [ - 0.5.0, - 0.5.3, - 0.6.1, 0.6.2, 0.7.0, 0.7.2, @@ -26,6 +23,7 @@ jobs: 0.9.1, 0.9.2, 0.10.0, + 0.10.1, ] # 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 diff --git a/requirements.txt b/requirements.txt index 671f5dd861..1127bc4e61 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.5.0, != 0.5.1, != 0.5.2, != 0.6.0, != 0.7.1, < 0.11.0 +jax >= 0.6.2, != 0.7.1, < 0.11.0 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8 diff --git a/tests/test_examples.py b/tests/test_examples.py index f47a7197c3..81d8f660c7 100644 --- a/tests/test_examples.py +++ b/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"): From 4bb2f93e587b80aba16f332802ad53704c8d9c24 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Fri, 29 May 2026 01:19:12 +0300 Subject: [PATCH 6/9] fix the failing test with jax> 0.10.1 due to minor difference in rounding for vmapped operation by looking at the eps range instead of exact range --- desc/integrals/_interp_utils.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index b77a32e7ea..40d9d35b0e 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -346,9 +346,14 @@ def polyroot_vec( if get_only_real_roots: a_min = -jnp.inf if a_min is None else a_min[..., jnp.newaxis] a_max = +jnp.inf if a_max is None else a_max[..., jnp.newaxis] + # Allow eps tolerance on the bounds so roots that land within a few ulp + # of a_min/a_max (e.g. a true root at x=0 returned as -1 ulp by the + # vectorized cubic formula) are not dropped. Clip survivors into range. r = jnp.where( - (jnp.abs(r.imag) <= eps) & (a_min <= r.real) & (r.real <= a_max), - r.real, + (jnp.abs(r.imag) <= eps) + & (a_min - eps <= r.real) + & (r.real <= a_max + eps), + jnp.clip(r.real, a_min, a_max), sentinel, ) From 6b7e871f1bffe9ad0249e0a4a5181e9453eb1f48 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sun, 31 May 2026 22:29:04 +0300 Subject: [PATCH 7/9] stick with jax<0.10.0 --- .github/workflows/jax_tests.yml | 4 ---- desc/integrals/_interp_utils.py | 9 ++------- requirements.txt | 2 +- 3 files changed, 3 insertions(+), 12 deletions(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index e14bdd79e0..82dbe52da4 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -22,11 +22,7 @@ jobs: 0.8.3, 0.9.1, 0.9.2, - 0.10.0, - 0.10.1, ] - # 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.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 diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index 40d9d35b0e..b77a32e7ea 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -346,14 +346,9 @@ def polyroot_vec( if get_only_real_roots: a_min = -jnp.inf if a_min is None else a_min[..., jnp.newaxis] a_max = +jnp.inf if a_max is None else a_max[..., jnp.newaxis] - # Allow eps tolerance on the bounds so roots that land within a few ulp - # of a_min/a_max (e.g. a true root at x=0 returned as -1 ulp by the - # vectorized cubic formula) are not dropped. Clip survivors into range. r = jnp.where( - (jnp.abs(r.imag) <= eps) - & (a_min - eps <= r.real) - & (r.real <= a_max + eps), - jnp.clip(r.real, a_min, a_max), + (jnp.abs(r.imag) <= eps) & (a_min <= r.real) & (r.real <= a_max), + r.real, sentinel, ) diff --git a/requirements.txt b/requirements.txt index 1127bc4e61..110d8633ef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, < 0.11.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.8 From 9d871370e287646fdd1cf11208f175fc6e8a5f58 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 1 Jun 2026 21:39:55 +0300 Subject: [PATCH 8/9] address Daniel's suggestions --- desc/particles.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/desc/particles.py b/desc/particles.py index 852522e96d..fff76a0029 100644 --- a/desc/particles.py +++ b/desc/particles.py @@ -415,10 +415,8 @@ def __init__( q=2, ): rho0, theta0, zeta0, xi0, E, m, q = map( - jnp.asarray, (rho0, theta0, zeta0, xi0, E, m, q) - ) - 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 @@ -550,8 +548,9 @@ def __init__( m=4, q=2, ): - R0, phi0, Z0, xi0, E, m, q = map(jnp.asarray, (R0, phi0, Z0, xi0, E, m, q)) - 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 @@ -694,8 +693,7 @@ def __init__( is_curve_magnetic_axis=False, ): self.curve = curve - E, m, q = map(jnp.asarray, (E, m, q)) - 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 @@ -854,8 +852,7 @@ def __init__( is_surface_from_eq=False, ): self.surface = surface - E, m, q = map(jnp.asarray, (E, m, q)) - 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 From 2dc5afe8e1df3ca1fc410a2bd13bfa5584837e0b Mon Sep 17 00:00:00 2001 From: YigitElma Date: Mon, 1 Jun 2026 21:41:14 +0300 Subject: [PATCH 9/9] remove jit in the same spirit as #2228 --- desc/particles.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/desc/particles.py b/desc/particles.py index fff76a0029..88e8359b93 100644 --- a/desc/particles.py +++ b/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 @@ def frame(self): """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.