JAX: add v0.4.32+

[adamjstewart]

• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.32
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.33
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.34
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.35
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.36
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.37
• https://github.com/jax-ml/jax/releases/tag/jax-v0.4.38

https://jax.readthedocs.io/en/latest/developer.html

[adamjstewart]

jax-ml/jax#23687
jax-ml/jax#23689

[wdconinc]

Checksums verified. Do you want to add py-jaxlib@=0.4.33 too now that it's out?

Deferring to maintainer on the py-jaxlib modifications, which are exercised in CI builds.

[adamjstewart]

jaxlib 0.4.33 hasn't been released on GitHub yet.

[wdconinc]

jaxlib 0.4.33 hasn't been released on GitHub yet.

You're right (of course). I mixed up jax and jaxlib.

[wdconinc]

I also see the CI build errors now.

[adamjstewart]

Yep. One of them is fixed in main, but not the other one. So I'm waiting to decide whether or not to backport the patch. JAX completely changed its entire build system to make it more difficult to use external CUDA and friends.

[aweits]

currently not building for me...

DEBUG: /scratch.local/spack7ucgbc8b/0811568bf239883f388d364b3e2f09d5/external/tsl/third_party/gpus/cuda/hermetic/cuda_configure.bzl:50:14: Cannot find None, either correct your path, or set the CLANG_CUDA_COMPILER_PATH or CC environment variables

build script for tsl looks to be missing some vars -- digging....

# Template values set by cuda_autoconf.
CPU_COMPILER = ('')
HOST_COMPILER_PATH = ('')

[adamjstewart]

Which compiler are you using?

[adamjstewart]

I have no idea how to fix the CPU installation and the GPU build is still quite broken. And I can't even build bazel on macOS, so I can't easily debug anything. So I'm definitely going to need help on this one. We either need to finish this or backport some complicated thing to get the aarch64 tests to pass, and I would rather finish this if possible.

[wdconinc]

This PR needs a rebase after #47260

[jonas-eschle]

Are there any news on this? It's holding up quite a few PRs, starting from TFP: #43688

[adamjstewart]

No solutions from upstream JAX maintainers, hoping @aweits or someone else can figure out how to get this to build because I can't.

[adamjstewart]

New bug: jax-ml/jax#25488

[adamjstewart]

Well, the build almost completes, but I can't even find an error message now.

[adamjstewart]

x86_64 build is working now, still need to fix the aarch64 build: jax-ml/jax#25598

[adamjstewart]

x86_64 CUDA build is still failing: jax-ml/jax#23689 (comment)