diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index 3736ab63b5..82dbe52da4 100644 --- a/.github/workflows/jax_tests.yml +++ b/.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 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..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. @@ -415,7 +414,8 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 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 diff --git a/requirements.txt b/requirements.txt index 80b282e29d..110d8633ef 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.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 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"):