Make map_coordinates differentiable for JAX 0.4.34

.github/workflows/jax_tests.yml

@@ -12,34 +12,35 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        jax-version: [0.3.0,  0.3.1,  0.3.2,  0.3.3,  0.3.4,  0.3.5,
-                      0.3.6,  0.3.7,  0.3.8,  0.3.9,  0.3.10, 0.3.11,
-                      0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.3.17,
-                      0.3.19, 0.3.20, 0.3.21, 0.3.22, 0.3.23, 0.3.24,
-                      0.3.25, 0.4.1,  0.4.2,  0.4.3,  0.4.4,  0.4.5,
-                      0.4.6,  0.4.7,  0.4.8,  0.4.9,  0.4.10, 0.4.11,
-                      0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18,
+        jax-version: [0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17,
                       0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23,
-                      0.4.24, 0.4.25]
+                      0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29,
+                      0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35]
+                      # 0.4.32 is not available on PyPI
+                      # earlier jax versions are not compatible with other
+                      # dependencies as of 2024-10-04
         group: [1, 2]
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.9
+      - name: Set up Python 3.10
         uses: actions/setup-python@v5
         with:
-          python-version: 3.9
-          cache: pip
-      - name: Install dependencies
+          python-version: '3.10'
+      - name: Upgrade pip
         run: |
           python -m pip install --upgrade pip
-          pip install -r devtools/dev-requirements.txt
-          pip install matplotlib==3.5.0
-      - name: Remove jax
+      - name: Install dependencies with given JAX version
         run: |
-          pip uninstall jax jaxlib -y
-      - name: install jax
+          sed -i '/jax/d' ./requirements.txt
+          sed -i '1i\jax[cpu] == ${{ matrix.jax-version }}' ./requirements.txt
+          cat ./requirements.txt
+          pip install -r ./devtools/dev-requirements.txt
+          pip install matplotlib==3.7.2
+      - name: Verify dependencies
         run: |
-          pip install "jax[cpu]==${{ matrix.jax-version }}"
+          python --version
+          pip --version
+          pip list
       - name: Test with pytest
         run: |
           pwd

CHANGELOG.md

@@ -24,6 +24,7 @@ Bug Fixes
 
 - Fixes bugs that occur when saving asymmetric equilibria as wout files
 - Fixes bug that occurs when using ``VMECIO.plot_vmec_comparison`` to compare to an asymmetric wout file
+- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version
 
 Deprecations
 

desc/backend.py

@@ -87,9 +87,12 @@
         treedef_is_leaf,
     )
 
-    trapezoid = (
-        jnp.trapezoid if hasattr(jnp, "trapezoid") else jax.scipy.integrate.trapezoid
-    )
+    if hasattr(jnp, "trapezoid"):
+        trapezoid = jnp.trapezoid  # for JAX 0.4.26 and later
+    elif hasattr(jax.scipy, "integrate"):
+        trapezoid = jax.scipy.integrate.trapezoid
+    else:
+        trapezoid = jnp.trapz  # for older versions of JAX, deprecated by jax 0.4.16
 
     def execute_on_cpu(func):
         """Decorator to set default device to CPU for a function.
@@ -200,6 +203,7 @@ def root_scalar(
         maxiter_ls=5,
         alpha=0.1,
         fixup=None,
+        full_output=False,
     ):
         """Find x where fun(x, *args) == 0.
 
@@ -227,6 +231,9 @@ def root_scalar(
         fixup : callable, optional
             Function to modify x after each update, ie to enforce periodicity. Should
             have a signature of the form fixup(x, *args) -> x'.
+        full_output : bool, optional
+            If True, also return a tuple where the first element is the residual from
+            the root finding and the second is the number of iterations.
 
         Returns
         -------
@@ -271,18 +278,25 @@ def bodyfun(state):
                 xk1, fk1 = backtrack(xk1, fk1, d)
                 return xk1, fk1, k1 + 1
 
-            state = guess, res(guess), 0
+            state = guess, res(guess), 0.0
             state = jax.lax.while_loop(condfun, bodyfun, state)
-            return state[0], state[1:]
+            if full_output:
+                return state[0], state[1:]
+            else:
+                return state[0]
 
         def tangent_solve(g, y):
             A = jax.jacfwd(g)(y)
             return y / A
 
-        x, (res, niter) = jax.lax.custom_root(
-            res, x0, solve, tangent_solve, has_aux=True
-        )
-        return x, (abs(res), niter)
+        if full_output:
+            x, (res, niter) = jax.lax.custom_root(
+                res, x0, solve, tangent_solve, has_aux=True
+            )
+            return x, (abs(res), niter)
+        else:
+            x = jax.lax.custom_root(res, x0, solve, tangent_solve, has_aux=False)
+            return x
 
     def root(
         fun,
@@ -294,6 +308,7 @@ def root(
         maxiter_ls=0,
         alpha=0.1,
         fixup=None,
+        full_output=False,
     ):
         """Find x where fun(x, *args) == 0.
 
@@ -321,6 +336,9 @@ def root(
         fixup : callable, optional
             Function to modify x after each update, ie to enforce periodicity. Should
             have a signature of the form fixup(x, *args) -> 1d array.
+        full_output : bool, optional
+            If True, also return a tuple where the first element is the residual from
+            the root finding and the second is the number of iterations.
 
         Returns
         -------
@@ -388,19 +406,26 @@ def bodyfun(state):
             state = (
                 jnp.atleast_1d(jnp.asarray(guess)),
                 jnp.atleast_1d(resfun(guess)),
-                0,
+                0.0,
             )
             state = jax.lax.while_loop(condfun, bodyfun, state)
-            return state[0], state[1:]
+            if full_output:
+                return state[0], state[1:]
+            else:
+                return state[0]
 
         def tangent_solve(g, y):
             A = jnp.atleast_2d(jax.jacfwd(g)(y))
             return _lstsq(A, jnp.atleast_1d(y))
 
-        x, (res, niter) = jax.lax.custom_root(
-            res, x0, solve, tangent_solve, has_aux=True
-        )
-        return x, (safenorm(res), niter)
+        if full_output:
+            x, (res, niter) = jax.lax.custom_root(
+                res, x0, solve, tangent_solve, has_aux=True
+            )
+            return x, (safenorm(res), niter)
+        else:
+            x = jax.lax.custom_root(res, x0, solve, tangent_solve, has_aux=False)
+            return x
 
 
 # we can't really test the numpy backend stuff in automated testing, so we ignore it
@@ -711,6 +736,7 @@ def root_scalar(
         maxiter_ls=5,
         alpha=0.1,
         fixup=None,
+        full_output=False,
     ):
         """Find x where fun(x, *args) == 0.
 
@@ -738,6 +764,9 @@ def root_scalar(
         fixup : callable, optional
             Function to modify x after each update, ie to enforce periodicity. Should
             have a signature of the form fixup(x) -> x'.
+        full_output : bool, optional
+            If True, also return a tuple where the first element is the residual from
+            the root finding and the second is the number of iterations.
 
         Returns
         -------
@@ -750,7 +779,10 @@ def root_scalar(
         out = scipy.optimize.root_scalar(
             fun, args, x0=x0, fprime=jac, xtol=tol, rtol=tol
         )
-        return out.root, out
+        if full_output:
+            return out.root, out
+        else:
+            return out.root
 
     def root(
         fun,
@@ -762,6 +794,7 @@ def root(
         maxiter_ls=0,
         alpha=0.1,
         fixup=None,
+        full_output=False,
     ):
         """Find x where fun(x, *args) == 0.
 
@@ -789,6 +822,9 @@ def root(
         fixup : callable, optional
             Function to modify x after each update, ie to enforce periodicity. Should
             have a signature of the form fixup(x, *args) -> 1d array.
+        full_output : bool, optional
+            If True, also return a tuple where the first element is the residual from
+            the root finding and the second is the number of iterations.
 
         Returns
         -------
@@ -803,7 +839,10 @@ def root(
         will solve it in a least squares sense.
         """
         out = scipy.optimize.root(fun, x0, args, jac=jac, tol=tol)
-        return out.x, out
+        if full_output:
+            return out.x, out
+        else:
+            return out.x
 
     def flatnonzero(a, size=None, fill_value=0):
         """A numpy implementation of jnp.flatnonzero."""

desc/basis.py

@@ -1573,7 +1573,7 @@ def zernike_radial(r, l, m, dr=0):
             "Analytic radial derivatives of Zernike polynomials for order>4 "
             + "have not been implemented."
         )
-    return s * jnp.where((l - m) % 2 == 0, out, 0)
+    return s * jnp.where((l - m) % 2 == 0, out, 0.0)
 
 
 def power_coeffs(l):
@@ -1732,7 +1732,7 @@ def _binom_body_fun(i, b_n):
     return b
 
 
-@custom_jvp
+@functools.partial(custom_jvp, nondiff_argnums=(4,))
 @jit
 @jnp.vectorize
 def _jacobi(n, alpha, beta, x, dx=0):
@@ -1804,13 +1804,13 @@ def _jacobi_body_fun(kk, d_p_a_b_x):
 
 
 @_jacobi.defjvp
-def _jacobi_jvp(x, xdot):
-    (n, alpha, beta, x, dx) = x
-    (ndot, alphadot, betadot, xdot, dxdot) = xdot
+def _jacobi_jvp(dx, x, xdot):
+    (n, alpha, beta, x) = x
+    (*_, xdot) = xdot
     f = _jacobi(n, alpha, beta, x, dx)
     df = _jacobi(n, alpha, beta, x, dx + 1)
     # in theory n, alpha, beta, dx aren't differentiable (they're integers)
     # but marking them as non-diff argnums seems to cause escaped tracer values.
     # probably a more elegant fix, but just setting those derivatives to zero seems
     # to work fine.
-    return f, df * xdot + 0 * ndot + 0 * alphadot + 0 * betadot + 0 * dxdot
+    return f, df * xdot

desc/equilibrium/coords.py

@@ -208,14 +208,18 @@ def fixup(y, *args):
                 fixup=fixup,
                 tol=tol,
                 maxiter=maxiter,
+                full_output=full_output,
                 **kwargs,
             )
         )
     )
     # See description here
     # https://github.com/PlasmaControl/DESC/pull/504#discussion_r1194172532
     # except we make sure properly handle periodic coordinates.
-    yk, (res, niter) = vecroot(yk, coords)
+    if full_output:
+        yk, (res, niter) = vecroot(yk, coords)
+    else:
+        yk = vecroot(yk, coords)
 
     out = compute(yk, outbasis)
     if full_output:
@@ -363,18 +367,28 @@ def fixup(x, *args):
                 fixup=fixup,
                 tol=tol,
                 maxiter=maxiter,
+                full_output=full_output,
                 **kwargs,
             )
         )
     )
     rho, theta_PEST, zeta = coords.T
-    theta, (res, niter) = vecroot(
-        # Assume λ=0 for default initial guess.
-        setdefault(guess, theta_PEST),
-        theta_PEST,
-        rho,
-        zeta,
-    )
+    if full_output:
+        theta, (res, niter) = vecroot(
+            # Assume λ=0 for default initial guess.
+            setdefault(guess, theta_PEST),
+            theta_PEST,
+            rho,
+            zeta,
+        )
+    else:
+        theta = vecroot(
+            # Assume λ=0 for default initial guess.
+            setdefault(guess, theta_PEST),
+            theta_PEST,
+            rho,
+            zeta,
+        )
     out = jnp.column_stack([rho, jnp.atleast_1d(theta.squeeze()), zeta])
     if full_output:
         return out, (res, niter)
@@ -466,6 +480,7 @@ def fixup(x, *args):
                 fixup=fixup,
                 tol=tol,
                 maxiter=maxiter,
+                full_output=full_output,
                 **kwargs,
             )
         )
@@ -474,7 +489,10 @@ def fixup(x, *args):
     if guess is None:
         # Assume λ=0 for default initial guess.
         guess = alpha + iota * zeta
-    theta, (res, niter) = vecroot(guess, alpha, rho, zeta, iota)
+    if full_output:
+        theta, (res, niter) = vecroot(guess, alpha, rho, zeta, iota)
+    else:
+        theta = vecroot(guess, alpha, rho, zeta, iota)
 
     out = jnp.column_stack([rho, jnp.atleast_1d(theta.squeeze()), zeta])
     if full_output:

desc/geometry/surface.py

@@ -741,10 +741,19 @@ def fun_jax(zeta_hat, theta, zeta):
             n, r, r_offset = n_and_r_jax(nodes)
             return jnp.arctan(r_offset[0, 1] / r_offset[0, 0]) - zeta
 
-        vecroot = jit(vmap(lambda x0, *p: root_scalar(fun_jax, x0, jac=None, args=p)))
-        zetas, (res, niter) = vecroot(
-            grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2]
+        vecroot = jit(
+            vmap(
+                lambda x0, *p: root_scalar(
+                    fun_jax, x0, jac=None, args=p, full_output=full_output
+                )
+            )
         )
+        if full_output:
+            zetas, (res, niter) = vecroot(
+                grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2]
+            )
+        else:
+            zetas = vecroot(grid.nodes[:, 2], grid.nodes[:, 1], grid.nodes[:, 2])
 
         zetas = np.asarray(zetas)
         nodes = np.vstack((np.ones_like(grid.nodes[:, 1]), grid.nodes[:, 1], zetas)).T

devtools/dev-requirements_conda.yml

@@ -15,9 +15,9 @@ dependencies:
   - pip:
     # Conda only parses a single list of pip requirements.
     # If two pip lists are given, all but the last list is skipped.
+    - jax >= 0.4.24, < 0.5.0
     - diffrax >= 0.4.1
     - interpax >= 0.3.3
-    - jax[cpu] >= 0.3.2, < 0.5.0
     - nvgpu
     - orthax
     - plotly >= 5.16, < 6.0
@@ -29,7 +29,6 @@ dependencies:
     - qicna @ git+https://github.com/rogeriojorge/pyQIC/
     - black[jupyter] = 24.3.0
 
-
   # building the docs
   - nbsphinx == 0.8.12
   - pandoc