Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group across 1 directory by dependabot[bot] · Pull Request #2235 · PlasmaControl/DESC

dependabot · 2026-06-03T19:20:10Z

⚠️ Dependabot is rebasing this PR ⚠️

Rebasing might not happen immediately, so don't worry if this takes some time.

Note: if you make any changes to this PR yourself, they will take precedence over the rebase.

Updates the requirements on jax to permit the latest version.
Updates jax to 0.10.1

Release notes

JAX v0.10.1

New features

Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098).

Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144).

Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144).

Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144).

Added jax.scipy.linalg.leslie for constructing Leslie matrices (#10144).

Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients (#10144).

Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices (#10144).

Added jax.scipy.linalg.helmert for constructing Helmert matrices (#10144).

Moved RNG APIs from "implementations" to dtypes (#27854):

Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.

jax.random.key and wrap_key_data now accept a dtype argument.

Breaking changes

with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

Deprecations

Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.

Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

Changelog

Sourced from jax's changelog.

JAX 0.10.1 (May 20, 2026)

New features

Added ResizeMethod.AREA to {func}jax.image.resize, which matches TensorFlow's AREA resizing ({jax-issue}[#20098](https://github.com/jax-ml/jax/issues/20098)).

Added {func}jax.scipy.linalg.hadamard for constructing Hadamard matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.circulant for constructing circulant matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.dft for constructing discrete Fourier transform matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.leslie for constructing Leslie matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.linalg.helmert for constructing Helmert matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).

Added {func}jax.scipy.special.boxcox and {func}jax.scipy.special.boxcox1p for the Box-Cox power transformation.

Moved RNG APIs from "implementations" to dtypes ({jax-issue}[#27854](https://github.com/jax-ml/jax/issues/27854)):

Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.

jax.random.key and wrap_key_data now accept a dtype argument.

Breaking changes

with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

Deprecations

Passing the copy, order, and ndmin arguments to {func}jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.

Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

JAX 0.10.0 (April 16, 2026)

New features:

Added ResizeMethod.CUBIC_PYTORCH to {func}jax.image.resize to match PyTorch's bicubic resize ({jax-issue}[#15768](https://github.com/jax-ml/jax/issues/15768)).

We now support differentiation of {func}jax.lax.linalg.qr for wide matrices and when full_matrices is True.

LAPACK operations are now parallelized along the batch dimension on CPU.

Added perturb_singular argument to {func}jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near-zero pivots in the LU decomposition. This is useful for

... (truncated)

Commits

619764c Apply patch to revert XLA PR #42396 on release/0.10.1 branch
433d98d Prepare for JAX release 0.10.1
8c38c32 [Pallas] Enable disabled TPU Interpret Mode test of example kernel
6820cf6 [pallas:sc] pl.kernel is now the only API for writing Pallas SC kernels
456b22d [state] Fixed a latent bug in _call_primitive_discharge_rule
d154b80 Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b...
11af068 Update XLA dependency to use revision http://github.com/openxla/xla/commit/31...
1ecdea1 Improve compilation throughput in multi-threaded auto-tuning environments. Th...
46f75a2 Make JaxprEqnContext a bit easier to work with making it easier to add new co...
7c01f9b [Pallas MPMD] Refactor tracing context
Additional commits viewable in compare view

github-actions · 2026-06-03T19:43:37Z

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    2.64 %    |     3.994e+03      |     4.099e+03      |    105.39    |       30.34        |       29.23        |
  test_proximal_jac_w7x_with_eq_update   |    4.72 %    |     6.618e+03      |     6.930e+03      |    312.22    |       148.13       |       150.64       |
  test_proximal_freeb_jac                |    0.61 %    |     1.339e+04      |     1.347e+04      |    81.17     |       81.54        |       80.52        |
  test_proximal_freeb_jac_blocked        |    1.05 %    |     7.724e+03      |     7.805e+03      |    81.28     |       70.39        |       66.96        |
  test_proximal_freeb_jac_batched        |    1.68 %    |     7.601e+03      |     7.728e+03      |    127.71    |       69.38        |       66.35        |
  test_proximal_jac_ripple               |    4.28 %    |     3.466e+03      |     3.614e+03      |    148.42    |       53.84        |       52.53        |
  test_proximal_jac_ripple_bounce1d      |    4.18 %    |     3.794e+03      |     3.952e+03      |    158.43    |       67.00        |       65.99        |
  test_eq_solve                          |    6.60 %    |     1.954e+03      |     2.083e+03      |    129.03    |       53.56        |       55.77        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

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](jax-ml/jax@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] <support@github.com>

YigitElma · 2026-06-19T08:47:17Z

Should we go back to updating by patch release? 0.10.1 was able to run, but in 0.10.2 some of the internal functions are removed and our jacfwd_chunked etc cannot be imported properly.

YigitElma · 2026-06-19T08:49:43Z

+        xtol=1e-16,
+        maxiter=200,
+        verbose=3,
+        options={"tr_method": "svd"},


Consider explaining this issue in tr_method docs for QR. If there is a chance that the jacobian is singular, then SVD is the more reliable option. If the problem is well-posed otherwise, having a wide bounds can cause this (some of the columns are completely 0).

github-actions · 2026-06-19T23:42:03Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     +1.44 +/- 2.78     | +1.24e-02 +/- 2.39e-02 |  8.73e-01 +/- 2.1e-02  |  8.61e-01 +/- 1.2e-02  |
 test_build_transform_fft_highres        |     +1.11 +/- 2.55     | +1.28e-02 +/- 2.95e-02 |  1.17e+00 +/- 2.1e-02  |  1.16e+00 +/- 2.1e-02  |
 test_equilibrium_init_lowres            |     +2.54 +/- 3.29     | +1.59e-01 +/- 2.06e-01 |  6.41e+00 +/- 1.5e-01  |  6.26e+00 +/- 1.4e-01  |
 test_objective_compile_atf              |     -2.17 +/- 3.99     | -1.31e-01 +/- 2.41e-01 |  5.92e+00 +/- 2.3e-01  |  6.05e+00 +/- 8.2e-02  |
 test_objective_compute_atf              |     +2.80 +/- 13.31    | +5.87e-05 +/- 2.79e-04 |  2.15e-03 +/- 1.6e-04  |  2.10e-03 +/- 2.3e-04  |
 test_objective_jac_atf                  |     +1.64 +/- 7.51     | +2.46e-02 +/- 1.13e-01 |  1.52e+00 +/- 9.3e-02  |  1.50e+00 +/- 6.3e-02  |
 test_perturb_1                          |     +0.22 +/- 1.68     | +3.33e-02 +/- 2.58e-01 |  1.54e+01 +/- 2.2e-01  |  1.54e+01 +/- 1.3e-01  |
 test_proximal_jac_atf                   |     +0.02 +/- 1.36     | +1.02e-03 +/- 6.95e-02 |  5.12e+00 +/- 3.9e-02  |  5.12e+00 +/- 5.8e-02  |
 test_proximal_freeb_compute             |     -4.66 +/- 3.60     | -7.60e-03 +/- 5.87e-03 |  1.56e-01 +/- 4.0e-03  |  1.63e-01 +/- 4.3e-03  |
 test_solve_fixed_iter                   |     -0.57 +/- 2.18     | -1.36e-01 +/- 5.17e-01 |  2.36e+01 +/- 4.3e-01  |  2.38e+01 +/- 2.8e-01  |
 test_objective_compute_ripple           |     -2.02 +/- 4.53     | -4.60e-03 +/- 1.03e-02 |  2.24e-01 +/- 6.7e-03  |  2.28e-01 +/- 7.8e-03  |
 test_objective_grad_ripple              |     -2.35 +/- 4.02     | -2.04e-02 +/- 3.48e-02 |  8.46e-01 +/- 2.0e-02  |  8.66e-01 +/- 2.9e-02  |
 test_build_transform_fft_lowres         |     +3.69 +/- 13.55    | +2.01e-02 +/- 7.38e-02 |  5.65e-01 +/- 1.1e-02  |  5.45e-01 +/- 7.3e-02  |
 test_equilibrium_init_medres            |     +3.52 +/- 3.04     | +1.58e-01 +/- 1.37e-01 |  4.65e+00 +/- 1.2e-01  |  4.50e+00 +/- 6.5e-02  |
 test_equilibrium_init_highres           |     +2.52 +/- 3.08     | +1.29e-01 +/- 1.58e-01 |  5.26e+00 +/- 1.1e-01  |  5.13e+00 +/- 1.2e-01  |
 test_objective_compile_dshape_current   |     -1.49 +/- 9.23     | -3.80e-02 +/- 2.36e-01 |  2.52e+00 +/- 2.3e-01  |  2.55e+00 +/- 3.2e-02  |
 test_objective_compute_dshape_current   |     -1.49 +/- 6.89     | -8.56e-06 +/- 3.96e-05 |  5.66e-04 +/- 2.1e-05  |  5.74e-04 +/- 3.3e-05  |
 test_objective_jac_dshape_current       |     -4.10 +/- 10.86    | -7.01e-04 +/- 1.85e-03 |  1.64e-02 +/- 9.8e-04  |  1.71e-02 +/- 1.6e-03  |
 test_perturb_2                          |     +0.47 +/- 1.41     | +5.77e-02 +/- 1.74e-01 |  1.23e+01 +/- 1.2e-01  |  1.23e+01 +/- 1.2e-01  |
 test_proximal_jac_atf_with_eq_update    |     -1.62 +/- 1.98     | -1.61e-01 +/- 1.97e-01 |  9.80e+00 +/- 1.1e-01  |  9.96e+00 +/- 1.7e-01  |
 test_proximal_freeb_jac                 |     +0.37 +/- 2.42     | +1.72e-02 +/- 1.13e-01 |  4.69e+00 +/- 1.0e-01  |  4.67e+00 +/- 5.1e-02  |
+test_solve_fixed_iter_compiled          |     -9.97 +/- 1.51     | -5.70e-01 +/- 8.61e-02 |  5.15e+00 +/- 7.4e-02  |  5.72e+00 +/- 4.4e-02  |
 test_LinearConstraintProjection_build   |     +1.80 +/- 3.39     | +7.87e-02 +/- 1.48e-01 |  4.45e+00 +/- 1.1e-01  |  4.37e+00 +/- 9.6e-02  |
 test_objective_compute_ripple_bounce1d  |     +0.37 +/- 11.96    | +7.19e-04 +/- 2.34e-02 |  1.96e-01 +/- 1.9e-02  |  1.96e-01 +/- 1.4e-02  |
 test_objective_grad_ripple_bounce1d     |     +2.77 +/- 4.76     | +2.32e-02 +/- 3.98e-02 |  8.58e-01 +/- 3.4e-02  |  8.35e-01 +/- 2.1e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

codecov · 2026-07-14T17:06:58Z

Codecov Report

❌ Patch coverage is 86.66667% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.33%. Comparing base (24aa7b9) to head (bc2ca4d).

Files with missing lines	Patch %	Lines
desc/batching.py	86.66%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2235   +/-   ##
=======================================
  Coverage   94.32%   94.33%           
=======================================
  Files         101      101           
  Lines       28866    28873    +7     
=======================================
+ Hits        27228    27236    +8     
+ Misses       1638     1637    -1

Files with missing lines	Coverage Δ
desc/batching.py	`83.13% <86.66%> (+0.71%)`	⬆️

... and 6 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

unalmis

Forced to block due to issue in above comment, and previous discussion.

unalmis · 2026-07-14T19:00:37Z

        return z1[mask], z2[mask]

    @pytest.mark.unit
+    @pytest.mark.xfail(reason="will be fixed by #2199")


#2222 (review).

It's the same problem as the above link, except this PR's approach is to break the code by releasing with a package that does not pass correctness tests.

We discussed this in the dev meeting, and the decision was to add this flag temporarily. Small numerical differences in different versions of dependencies can already be causing problems that we haven't seen. For example, this workflow fails due to jax-finufft=1.2.0, currently this workflow is the only one that uses that version because test_compute_everything happens to run on the runner that uses Python=3.10. What I mean is we don't test each dependency Python version combination on CI, and if certain parts of the code are too sensitive to these, there's not much to do.

There won't be a release very soon. In the meantime, we can merge the proper fix #2199. I can approve that PR and we can merge it soon, but it would be a lot easier if you spent a couple of minutes to cherry-pick a couple of commits to point it to master. Otherwise, it will wait on a chain of PRs. I can say the same for couple other PRs too. You can keep the original PRs open, I am pretty sure if you merge the exact same changes to master, there won't be merge conflicts.

I worked on DESC for 4 years. I will block PRs that 1) show failure in the correctness of the code, and 2) remove required testing infrastructure as means to silence such failures.

The other workflow things are non-sequitur.

Nothing is holding up the PRs you mentioned. My existing PRs that allow the code to work in a broader range of configurations have been ready for merge with no issues.

In the meantime, we can merge the proper fix #2199.

You need #2170 and #2199 ; reasons why have been discussed in prs. The ordering of the PR is deliberate for correctness.

dependabot Bot added dependencies Issue related to libraries we depend on and how we interface with them run_benchmarks Run timing benchmarks on this PR against current master branch skip_changelog No need to update changelog on this PR labels Jun 3, 2026

YigitElma mentioned this pull request Jun 3, 2026

Workaround for dependabot issue #2234

Merged

YigitElma marked this pull request as draft June 3, 2026 22:23

dependabot Bot changed the title ~~Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group~~ Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group across 1 directory Jun 9, 2026

dependabot Bot force-pushed the dependabot/pip/jax-93493677ae branch from 3d11039 to a71b715 Compare June 9, 2026 17:23

YigitElma mentioned this pull request Jun 17, 2026

Use qr_multiply to allow better reuse of factorisation #2244

Merged

YigitElma and others added 2 commits June 17, 2026 19:10

use SVD for possibly rand-deficient system

9adb60f

Merge branch 'master' into dependabot/pip/jax-93493677ae

05b3097

YigitElma mentioned this pull request Jun 19, 2026

Run benchmarks for dependency updates #2246

Closed

YigitElma reviewed Jun 19, 2026

View reviewed changes

YigitElma and others added 5 commits June 19, 2026 13:03

Merge branch 'master' into dependabot/pip/jax-93493677ae

a285c98

trying to solve removed internal functions, credit to Claude

69acac3

update dependecy ranges

4c32861

fix deps

50477a2

temporarily remove 0.10.2

8c4a708

YigitElma and others added 4 commits July 9, 2026 18:09

Merge branch 'master' into dependabot/pip/jax-93493677ae

e76b870

add xfail to test, try 0.10.2

89abded

Merge branch 'master' into dependabot/pip/jax-93493677ae

9ff1025

remove 0.10.2 due to bug in jax-finufft

bc2ca4d

YigitElma marked this pull request as ready for review July 14, 2026 18:08

unalmis requested changes Jul 14, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group across 1 directory#2235

Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group across 1 directory#2235
dependabot[bot] wants to merge 12 commits into
masterfrom
dependabot/pip/jax-93493677ae

dependabot Bot commented on behalf of github Jun 3, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 3, 2026 •

edited

Loading

Uh oh!

YigitElma commented Jun 19, 2026

Uh oh!

YigitElma Jun 19, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 19, 2026 •

edited

Loading

Uh oh!

codecov Bot commented Jul 14, 2026

Uh oh!

unalmis left a comment

Uh oh!

unalmis Jul 14, 2026

Uh oh!

YigitElma Jul 14, 2026

Uh oh!

unalmis Jul 15, 2026 •

edited

Loading

Uh oh!

unalmis Jul 15, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dependabot Bot commented on behalf of github Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

JAX v0.10.1

JAX 0.10.1 (May 20, 2026)

JAX 0.10.0 (April 16, 2026)

Uh oh!

github-actions Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Memory benchmark result

Uh oh!

YigitElma commented Jun 19, 2026

Uh oh!

YigitElma Jun 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Jul 14, 2026

Codecov Report

Uh oh!

unalmis left a comment

Choose a reason for hiding this comment

Uh oh!

unalmis Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

YigitElma Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

unalmis Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

unalmis Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dependabot Bot commented on behalf of github Jun 3, 2026 •

edited

Loading

github-actions Bot commented Jun 3, 2026 •

edited

Loading

YigitElma Jun 19, 2026 •

edited

Loading

github-actions Bot commented Jun 19, 2026 •

edited

Loading

unalmis Jul 15, 2026 •

edited

Loading

unalmis Jul 15, 2026 •

edited

Loading