python312Packages.{jax,jaxlib}: 0.4.28 -> 0.4.38; use jaxlib-bin by default #363130
+391
−205
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
natsukium
commented
Dec 8, 2024
|
GaetanLepage
requested changes
Dec 8, 2024
| @@ -168,10 +174,10 @@ buildPythonPackage { | |||
|
|
|||
| disabled = | |||
There was a problem hiding this comment.
Suggested change
| disabled = | |
| disabled = pythonOlder "3.10"; |
I am not sure that keeping an upper bound makes a ton of sense.
Actually, I would be in favor of removing the disabled attribute entirely.
There was a problem hiding this comment.
I'm not aggressively trying to remove it, but it seems a bit redundant here.
Comment on lines
240
to
179
| platforms = [ | ||
| "aarch64-linux" | ||
| "aarch64-darwin" | ||
| "x86_64-linux" | ||
| "x86_64-darwin" | ||
| ]; |
There was a problem hiding this comment.
Suggested change
| platforms = [ | |
| "aarch64-linux" | |
| "aarch64-darwin" | |
| "x86_64-linux" | |
| "x86_64-darwin" | |
| ]; |
Shouldn't we remove the platforms attribute entirely as they are now all supported (which is the default) ?
Comment on lines
48
to
61
| # jaxlib is _not_ included in propagatedBuildInputs because there are | ||
| # different versions of jaxlib depending on the desired target hardware. The | ||
| # JAX project ships separate wheels for CPU, GPU, and TPU. | ||
| propagatedBuildInputs = [ | ||
| dependencies = [ | ||
| jaxlib |
There was a problem hiding this comment.
Is this comment and #156808 no longer relevant now that jax depends on jaxlib directly with the CUDA plugin change in 0.4.30? I was sitting on a PR removing jaxlib from {dm-haiku, equinox, flax, optax, etc.} but I guess it's no longer relevant. And should packages depending on jax also put jaxlib in their dependencies after this PR or is that no longer necessary as jax itself now depends on jaxlib?
Thanks for tackling this by the way!
There was a problem hiding this comment.
Is this comment and #156808 no longer relevant now that jax depends on jaxlib directly with the CUDA plugin change in 0.4.30?
That is also my understanding
I was sitting on a PR removing jaxlib from {dm-haiku, equinox, flax, optax, etc.} but I guess it's no longer relevant. And should packages depending on jax also put jaxlib in their dependencies after this PR or is that no longer necessary as jax itself now depends on jaxlib?
I agree that this situation is a bit confusing currently. I propose that we merge this PR along with written guidance to maintainers on migrating.
There was a problem hiding this comment.
Yes, from my understanding, we will be able to get rid of specifying jaxlib in downstream packages.
There was a problem hiding this comment.
I'll delete this comment as it causes confusion.
I propose that we merge this PR along with written guidance to maintainers on migrating.
Do you mean to add in the release notes that it is not necessary to explicitly add jaxlib as a dependency for packages that depend on jax?
There was a problem hiding this comment.
Do you mean to add in the release notes that it is not necessary to explicitly add jaxlib as a dependency for packages that depend on jax?
Precisely
There was a problem hiding this comment.
I added the information to the release note.
b69c1a4
|
Trying to "build" |
Looks like we can just remove it from EDIT: bingo I think we can safely get rid of it now |
|
@natsukium ping. |
|
0.4.38 is out: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.38 |
|
I had been away from home for a while on business, and I was sick in bed. |
Oh sorry to hear that. I hope that you will feel better soon :) |
|
We're about to merge staging-next today or tomorrow I believe. |
|
GaetanLepage
force-pushed
the
jax-jaxlib/fix
branch
2 times, most recently
from
61a3070 to
6cec536
Compare
GaetanLepage
force-pushed
the
jax-jaxlib/fix
branch
2 times, most recently
from
1ef503c to
9f1ecfe
Compare
GaetanLepage
changed the title
|
|
|
GaetanLepage
force-pushed
the
jax-jaxlib/fix
branch
from
9f1ecfe to
ac1eb00
Compare
|
Sorry for the slow communication. I'm back. |
13 tasks
Sure ! Thanks for your effort @natsukium :) |
|
Apparently, evaluation fails. |
github-actions
bot
added
6.topic: nixos
natsukium
force-pushed
the
jax-jaxlib/fix
branch
from
ccd59d3 to
55b9b26
Compare
Thanks to bazel, it is difficult to source build with our current resources.
natsukium
force-pushed
the
jax-jaxlib/fix
branch
2 times, most recently
from
62b4b21 to
9d7f862
Compare
|
I think it would be nice to merge #368790 first in order to properly test this PR. |
natsukium
force-pushed
the
jax-jaxlib/fix
branch
from
9d7f862 to
b070bad
Compare
|
I think it is ready for the merge. |
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
python312Packages.jaxlib: use jaxlib-bin by default
Thanks to bazel, it is difficult to source build with our current resources.
python312Packages.jaxlib-bin: 0.4.28 -> 0.4.36
Diff: jax-ml/jax@jaxlib-v0.4.28...jaxlib-v0.4.38
Changelog: https://github.com/jax-ml/jax/releases/tag/jaxlib-v0.4.38
python312Packages.jax: 0.4.28 -> 0.4.36
Diff: jax-ml/jax@jax-v0.4.28...jax-v0.4.38
Changelog: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.38
Add binaries for aarch64-linux
Things done
nix.conf? (See Nix manual)sandbox = relaxedsandbox = truenix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD". Note: all changes have to be committed, also see nixpkgs-review usage./result/bin/)Add a 👍 reaction to pull requests you find important.