python39Packages.jax build failure on x86_64-linux as of ed65ee50

[samuela]

Issue description

Build of python39Packages.jax failed on x86_64-linux as of ed65ee5. This is currently breaking python39Packages.optax.

error: builder for '/nix/store/hrvisx8ps12yv6ws8di9ppmbygqsz8jm-python3.9-jax-0.4.1.drv' failed with exit code 1;
       last 10 log lines:
       > =========================== short test summary info ============================
       > ERROR tests/ode_test.py
       > ERROR tests/scipy_fft_test.py
       > ERROR tests/scipy_interpolate_test.py
       > ERROR tests/scipy_optimize_test.py
       > ERROR tests/scipy_signal_test.py
       > ERROR tests/scipy_stats_test.py
       > ERROR tests/third_party/scipy/line_search_test.py
       > ========== 10640 passed, 1085 skipped, 7 errors in 1662.15s (0:27:42) ==========
       > /nix/store/sw36plhp82916wwg6i6097rkzza7d950-stdenv-linux/setup: line 1594: pop_var_context: head of shell_variables not a function context
       For full logs, run 'nix log /nix/store/hrvisx8ps12yv6ws8di9ppmbygqsz8jm-python3.9-jax-0.4.1.drv'.
error: 1 dependencies of derivation '/nix/store/xp50bma92ibaw81zvmsj0ykzywhb1k7p-python3.9-chex-0.1.6.drv' failed to build
error: 1 dependencies of derivation '/nix/store/aih766h5kybk4pdf0rdiv5lk92hka5a4-python3.9-optax-0.1.4.drv' failed to build

full build log

python39Packages.optax maintainers cc: @ndl
python39Packages.jax maintainers cc: @samuela
Other cc: n/a

Steps to reproduce

1. Checkout nixpkgs at commit ed65ee5
2. Run nix-build -A python39Packages.jax

Technical details

 - system: `"x86_64-linux"`
 - host os: `Linux 5.15.0-1034-azure, Ubuntu, 22.04.2 LTS (Jammy Jellyfish), nobuild`
 - multi-user?: `yes`
 - sandbox: `yes`
 - version: `nix-env (Nix) 2.13.3`
 - nixpkgs: `/nix/store/qdp2s4gmwz57iz8f9pibglxibdzz6lil-source`

Misc.

This issue was automatically generated by nixpkgs-upkeep.

• CI workflow that created this issue.
• Internal tags: -5443535557983542332 -6718544890248038006

[samuela]

dup of #221739