From a71b715f74398ae80117eac5225498e6826a6b5c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 9 Jun 2026 17:23:26 +0000 Subject: [PATCH 1/8] Update jax requirement in the jax group across 1 directory Updates the requirements on [jax](https://github.com/jax-ml/jax) to permit the latest version. Updates `jax` to 0.10.1 - [Release notes](https://github.com/jax-ml/jax/releases) - [Changelog](https://github.com/jax-ml/jax/blob/main/CHANGELOG.md) - [Commits](https://github.com/jax-ml/jax/compare/jax-v0.6.2...jax-v0.10.1) --- updated-dependencies: - dependency-name: jax dependency-version: 0.10.1 dependency-type: direct:production dependency-group: jax ... Signed-off-by: dependabot[bot] --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 110d8633ef..1127bc4e61 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, < 0.10.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 From 9adb60f349bc5b4deb0f5f86d09ff2a66cfdfe1e Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 17 Jun 2026 19:10:30 +0300 Subject: [PATCH 2/8] use SVD for possibly rand-deficient system --- devtools/dev-requirements.txt | 2 +- tests/test_examples.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/devtools/dev-requirements.txt b/devtools/dev-requirements.txt index 07ad4887ef..e83a283120 100644 --- a/devtools/dev-requirements.txt +++ b/devtools/dev-requirements.txt @@ -32,7 +32,7 @@ pytest ~= 9.0 pytest-benchmark <= 5.2.3 pytest-cov >= 2.6.0, <= 7.1.0 pytest-monitor <= 1.6.6 -pytest-mpl == 0.16.1 +pytest-mpl >= 0.16.1, <= 0.19.0 pytest-split >= 0.8.2, <= 0.11.0 qicna @ git+https://github.com/rogeriojorge/pyQIC/ qsc <= 0.1.3 diff --git a/tests/test_examples.py b/tests/test_examples.py index 81d8f660c7..f8fc4434a5 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -174,7 +174,14 @@ def test_solve_bounds(): obj = ObjectiveFunction( ForceBalance(normalize=False, normalize_target=False, bounds=(-3e3, 3e3), eq=eq) ) - eq.solve(objective=obj, ftol=1e-16, xtol=1e-16, maxiter=200, verbose=3) + eq.solve( + objective=obj, + ftol=1e-16, + xtol=1e-16, + maxiter=200, + verbose=3, + options={"tr_method": "svd"}, + ) # check that all errors are nearly 0, since residual values are within target bounds f = obj.compute_scaled_error(obj.x(eq)) From 69acac3bfd1d085127b67aab2325f8dfddbaee6f Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sat, 20 Jun 2026 01:31:04 +0300 Subject: [PATCH 3/8] trying to solve removed internal functions, credit to Claude --- desc/batching.py | 43 +++++++++++++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/desc/batching.py b/desc/batching.py index eec6bedb51..5ddee4b2b7 100644 --- a/desc/batching.py +++ b/desc/batching.py @@ -9,9 +9,7 @@ _check_output_dtype_jacrev, _jacfwd_unravel, _jacrev_unravel, - _jvp, _std_basis, - _vjp, ) from jax._src.api_util import _ensure_index, argnums_partial, check_callable from jax._src.numpy.vectorize import ( @@ -409,6 +407,27 @@ def wrapped(*args, **kwargs): return wrapped +def _argnums_partial(fun, argnums, args, kwargs): + """Bind every argument except those in ``argnums``. + + This mirrors JAX's internal ``argnums_partial`` but returns a plain callable + (and the tuple of differentiated arguments) instead of a + ``linear_util.WrappedFun``. A plain callable is what the public + ``jax.jvp``/``jax.vjp`` expect, so we avoid the private ``jax._src`` helpers + (e.g. ``_jvp``, which was removed in JAX 0.10.2). + """ + argnums_t = (argnums,) if isinstance(argnums, int) else tuple(argnums) + dyn_args = tuple(args[i] for i in argnums_t) + + def f_partial(*dyn): + full_args = list(args) + for i, a in zip(argnums_t, dyn): + full_args[i] = a + return fun(*full_args, **kwargs) + + return f_partial, dyn_args + + def jacfwd_chunked( fun, argnums=0, @@ -463,18 +482,17 @@ def jacfwd_chunked( @wraps(fun, docstr=docstr, argnums=argnums) def jacfun(*args, **kwargs): - f = lu.wrap_init(fun, kwargs) - f_partial, dyn_args = argnums_partial( - f, argnums, args, require_static_args_hashable=False - ) + f_partial, dyn_args = _argnums_partial(fun, argnums, args, kwargs) tree_map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args) if not has_aux: - pushfwd = partial(_jvp, f_partial, dyn_args) + pushfwd = lambda tangents: jax.jvp(f_partial, dyn_args, tangents) y, jac = vmap_chunked(pushfwd, chunk_size=chunk_size)(_std_basis(dyn_args)) y = tree_map(lambda x: x[0], y) jac = tree_map(lambda x: jnp.moveaxis(x, 0, -1), jac) else: - pushfwd = partial(_jvp, f_partial, dyn_args, has_aux=True) + pushfwd = lambda tangents: jax.jvp( + f_partial, dyn_args, tangents, has_aux=True + ) y, jac, aux = vmap_chunked(pushfwd, chunk_size=chunk_size)( _std_basis(dyn_args) ) @@ -550,15 +568,12 @@ def jacrev_chunked( @wraps(fun, docstr=docstr, argnums=argnums) def jacfun(*args, **kwargs): - f = lu.wrap_init(fun, kwargs) - f_partial, dyn_args = argnums_partial( - f, argnums, args, require_static_args_hashable=False - ) + f_partial, dyn_args = _argnums_partial(fun, argnums, args, kwargs) tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args) if not has_aux: - y, pullback = _vjp(f_partial, *dyn_args) + y, pullback = jax.vjp(f_partial, *dyn_args) else: - y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True) + y, pullback, aux = jax.vjp(f_partial, *dyn_args, has_aux=True) tree_map(partial(_check_output_dtype_jacrev, holomorphic), y) jac = vmap_chunked(pullback, chunk_size=chunk_size)(_std_basis(y)) jac = jac[0] if isinstance(argnums, int) else jac From 4c32861e6798754ab9f05ce3f0c36f5a5cbdbe35 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sat, 20 Jun 2026 01:32:27 +0300 Subject: [PATCH 4/8] update dependecy ranges --- .github/workflows/jax_tests.yml | 3 +++ requirements.txt | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/jax_tests.yml b/.github/workflows/jax_tests.yml index 82dbe52da4..b89cd07ed6 100644 --- a/.github/workflows/jax_tests.yml +++ b/.github/workflows/jax_tests.yml @@ -22,6 +22,9 @@ jobs: 0.8.3, 0.9.1, 0.9.2, + 0.10.0, + 0.10.1, + 0.10.2, ] # 0.7.0 have performance issues but we still support it, see diffrax#680 # 0.7.1 fails with equinox, see equinox#1081 diff --git a/requirements.txt b/requirements.txt index 1127bc4e61..f25cf83f4e 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.2 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8 From 50477a2e374a3872432920dc30c64955718be677 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sat, 20 Jun 2026 01:34:34 +0300 Subject: [PATCH 5/8] fix deps --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index f25cf83f4e..24285f30b0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, =< 0.10.2 +jax >= 0.6.2, != 0.7.1, <= 0.10.2 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8 From 8c4a708db0c705c3e7bb7e979f68116a0fd61ecf Mon Sep 17 00:00:00 2001 From: YigitElma Date: Sat, 20 Jun 2026 02:08:18 +0300 Subject: [PATCH 6/8] temporarily remove 0.10.2 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 24285f30b0..d8ccc6c24e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, <= 0.10.2 +jax >= 0.6.2, != 0.7.1, <= 0.10.1 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8 From 89abdedba8f36a37ebef3afdcd0f991a1a87fd48 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 14 Jul 2026 10:23:33 -0400 Subject: [PATCH 7/8] add xfail to test, try 0.10.2 --- requirements.txt | 2 +- tests/test_integrals.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index a5eb30cbff..326e74bdeb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, <= 0.10.1 +jax >= 0.6.2, != 0.7.1, <= 0.10.2 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8 diff --git a/tests/test_integrals.py b/tests/test_integrals.py index d26761ec19..cd4255cefa 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -724,6 +724,7 @@ def filter(z1, z2): return z1[mask], z2[mask] @pytest.mark.unit + @pytest.mark.xfail(reason="will be fixed by #2199") def test_z1_first(self): """Case where straight line through first two intersects is in epigraph.""" start = np.pi / 3 From bc2ca4d18223428470a44d3988a95c1401bbd38d Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 14 Jul 2026 12:33:21 -0400 Subject: [PATCH 8/8] remove 0.10.2 due to bug in jax-finufft --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 326e74bdeb..a5eb30cbff 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -jax >= 0.6.2, != 0.7.1, <= 0.10.2 +jax >= 0.6.2, != 0.7.1, <= 0.10.1 colorama <= 0.4.6 diffrax >= 0.6.0, <= 0.7.2 equinox >=0.11.10, <=0.13.8