python310Packages.jax build failure on x86_64-linux as of 7554374d

[samuela]

Issue description

Build of python310Packages.jax failed on x86_64-linux as of 7554374. This is currently breaking python3Packages.jax.

error: builder for '/nix/store/cmfdzlv16zm0k0nbia6spgjr8ffbirp9-python3.10-jax-0.3.6.drv' failed with exit code 1;
       last 10 log lines:
       > FAILED tests/sparse_test.py::cuSparseTest::test_csr_fromdense_ad_float32[8,8]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_complex64[5,5]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_complex64[5,8]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_complex64[8,5]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_complex64[8,8]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_float32[5,5]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_float32[5,8]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_float32[8,5]
       > FAILED tests/sparse_test.py::cuSparseTest::test_coo_fromdense_ad_float32[8,8]
       > ========= 59 failed, 14190 passed, 1447 skipped in 1589.83s (0:26:29) ==========
       For full logs, run 'nix log /nix/store/cmfdzlv16zm0k0nbia6spgjr8ffbirp9-python3.10-jax-0.3.6.drv'.

full build log

python3Packages.jax maintainers cc: @samuela
python310Packages.jax maintainers cc: @samuela
Other cc: n/a

Steps to reproduce

1. Checkout nixpkgs at commit 7554374
2. Run nix-build -A python310Packages.jax

Technical details

 - system: `"x86_64-linux"`
 - host os: `Linux 5.15.0-1014-azure, Ubuntu, 20.04.4 LTS (Focal Fossa), nobuild`
 - multi-user?: `yes`
 - sandbox: `yes`
 - version: `nix-env (Nix) 2.10.3`
 - nixpkgs: `/nix/store/vg2xri6d14pzcqb0kh81fdaqv8fmkmfr-source`

Misc.

This issue was automatically generated by nixpkgs-upkeep.

• CI workflow that created this issue.
• Internal tags: 6226879539297373847 -4932070138951427960

[nixos-discourse]

This issue has been mentioned on NixOS Discourse. There might be relevant details there:

https://discourse.nixos.org/t/nixpkgss-current-development-workflow-is-not-sustainable/18741/99

[samuela]

latest release appears to be https://github.com/google/jax/releases/tag/jax-v0.3.15. I'll try upgrading jax/jaxlib to that to see if that fixes this. o/w we should file an upstream bug report

[mcwitt]

@samuela in case it's useful, I'd started to work on upgrading to 0.3.15 in #183051, but ran into issues I haven't been able to resolve with the darwin builds (linux build is working)

[mcwitt]

Unfortunately it seems there's a bit of a chicken-and-egg problem because v0.3.15 depends on etils, which in turn depends on jax. There's a PR adding etils (#183052), but we can't merge it without jax building (unless we're willing to merge it in a broken state).

[uri-canva]

@mcwitt for the darwin builds take a look at #184395, if you have enough context you can review it, if not you can pick out some of the changes and add them to your own PR and see if they help.

[uri-canva]

Oh sorry, just realised this is for jax, and my fix is for jaxlib, maybe you can find something useful but maybe not.

[samuela]

@mcwitt I'm working on a temporary fix to jax atm that just disables the failing tests (created upstream issue google/jax#11722). Then I say we merge the etils PR, and then attempt a jax/jaxlib upgrade and try un-disabling the tests.

[samuela]

Created #185078 to fix the current jax version. Hoping to get that merged by tomorrow.