python311Packages.orbax-checkpoint: init at 0.5.3 #266467
Conversation
GaetanLepage
force-pushed
the
orbax-checkpoint
branch
from
abcee51
to
f154716
Compare
|
Result of 4 packages built:
|
GaetanLepage
force-pushed
the
orbax-checkpoint
branch
4 times, most recently
from
6804158
to
4af94bc
Compare
| preCheck = '' | ||
| export HOME=$(mktemp -d) | ||
| ''; |
There was a problem hiding this comment.
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.
jaxlib should never be in propagatedBuildInputs. if setup.py asks for it, remove with pythonRemoveDeps
There was a problem hiding this comment.
is upgrading tensorstore required by orbax-checkpoint?
There was a problem hiding this comment.
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.
GaetanLepage
changed the title
GaetanLepage
force-pushed
the
orbax-checkpoint
branch
from
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.
Version 0.7.5 of flax fails to build because jax.config.define_*_state was removed in jax 0.4.25.
nix build .#python311Packages.flaxerror: 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.
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.
@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 = 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/)