-
-
Notifications
You must be signed in to change notification settings - Fork 14.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
python311Packages.orbax-checkpoint: init at 0.5.3 #266467
Conversation
abcee51
to
f154716
Compare
Result of 4 packages built:
|
6804158
to
4af94bc
Compare
preCheck = '' | ||
export HOME=$(mktemp -d) | ||
''; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why is this necessary? also why not put it in the overridePythonAttrs
below?
absl-py | ||
etils | ||
jax | ||
jaxlib |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
jaxlib
should never be in propagatedBuildInputs
. if setup.py asks for it, remove with pythonRemoveDeps
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is upgrading tensorstore required by orbax-checkpoint?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
After orbax-checkpoint version 0.4.3 I run into failures at runtime with the current version of tensorstore. I see despite the title that the orbax version is at 0.4.8 here.
4af94bc
to
ca76d94
Compare
@@ -21,7 +22,7 @@ | |||
, tensorflow | |||
}: | |||
|
|||
buildPythonPackage rec { | |||
let flax = buildPythonPackage rec { | |||
pname = "flax"; | |||
version = "0.7.5"; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Version 0.7.5 of flax fails to build because jax.config.define_*_state
was removed in jax 0.4.25.
nix build .#python311Packages.flax
error: builder for '/nix/store/w04njz8gwd8q2qa2giwlfg7yicb8jjhi-python3.11-flax-0.7.5.drv' failed with exit code 1;
last 10 log lines:
> File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
> File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/__init__.py", line 19, in <module>
> from .configurations import (
> File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/configurations.py", line 93, in <module>
> flax_filter_frames = define_bool_state(
> ^^^^^^^^^^^^^^^^^^
> File "/nix/store/f03b4mz4vmnl70awzsj84ywpdb37ksls-python3.11-flax-0.7.5/lib/python3.11/site-packages/flax/configurations.py", line 42, in define_bool_state
> return jax_config.define_bool_state('flax_' + name, default, help)
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> AttributeError: 'Config' object has no attribute 'define_bool_state'
For full logs, run 'nix-store -l /nix/store/w04njz8gwd8q2qa2giwlfg7yicb8jjhi-python3.11-flax-0.7.5.drv'.
Bumping to the latest version 0.8.2 fixes the issue (it was fixed in 0.8.0).
|
||
let orbax-checkpoint = buildPythonPackage rec { | ||
pname = "orbax-checkpoint"; | ||
version = "0.5.3"; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The latest version is 0.5.7.
@@ -8892,6 +8892,8 @@ self: super: with self; { | |||
|
|||
oras = callPackage ../development/python-modules/oras { }; | |||
|
|||
orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint { }; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@GaetanLepage perhaps out of scope for this PR, but orbax-checkpoint's dependency on protobuf causes a conflict with tensorflow since tensorflow has its own protobuf version. This means python packages can't depend on orbax-checkpoint and tensorflow at the same time. See tensorflow-build
in pkgs/top-level/python-packages.nix.
tensorflow-build = let
compat = rec {
protobufTF = pkgs.protobuf_21.override {
abseil-cpp = pkgs.abseil-cpp_202301;
};
...
I first noticed this when trying to package gpjax with the following code.
nix why-depends --all .#gpjax /nix/store/cb1gs888vfqxawvc65q1dk6jzbayh3wz-source#protobuf_21
/nix/store/icb03rmfzgph4bhxm4wic934qwx7y8j7-python3.11-gpjax-0.8.2
├───/nix/store/n8blzy5lqs9gl7yi4m8q5c44gnfkdgcc-python3.11-orbax-checkpoint-0.5.7
│ └───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
│ └───/nix/store/swx2q314if8j37ryn35nr2yb60axmrxz-protobuf-21.12
└───/nix/store/3smllqadx197wmspwd3zkyrxz2rrj01b-python3.11-tensorflow-probability-0.21.0
└───/nix/store/b6fvycgjpin2g72mfv5fw3f9agqhn697-python3.11-tensorflow-2.13.0
├───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
├───/nix/store/api4113pl9qg9296cf15igrfnxyifjhl-python3.11-tensorboard-2.16.2
│ ├───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
│ └───/nix/store/d4dyq3y1scsaprv1y72fvz7sbn31dqjb-python3.11-tensorboard_plugin_profile-2.11.1
│ └───/nix/store/vq7w8s4acdn9kx3638pz186ar0wsh5b9-python3.11-protobuf-4.21.12
└───/nix/store/d4dyq3y1scsaprv1y72fvz7sbn31dqjb-python3.11-tensorboard_plugin_profile-2.11.1
It doesn't seem tensorflow-build
exposes protobuf-pythonTF
ergonomically but
orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint { };
could still be replaced with
orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint {
protobuf = pkgs.protobuf_21.override {
abseil-cpp = pkgs.abseil-cpp_202301;
};
}
with the caveat that the two have to stay in sync. I'm a bit confused what the current situation is, for example, pkgs.python3Packages.tensorboard
also conflicts with tensorflow which seems strange to me. Why does tensorflow carry its own versions of packages that are also exposed at the top level?
@GaetanLepage sorry for the noise, but this PR seems to be superseded by |
i'll close for now since it looks like these changes have been superseded by other PRs, and we can always reopen or open follow up PRs to resurrect anything from this one |
Description of changes
Add orbax-checkpoint.
Things done
nix.conf
? (See Nix manual)sandbox = relaxed
sandbox = true
nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD"
. Note: all changes have to be committed, also see nixpkgs-review usage./result/bin/
)