gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.jax: 0.7.2 -> 0.8.0 (#452567)

Browse Source
This commit is contained in:
kirillrdy
2025-10-20 21:37:51 +00:00
committed by GitHub
parent 990146f6a3 c98a6d2187
commit 5d551d15ea
9 changed files with 63 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/dm-haiku/default.nix
+20 -19
View File
@@ -36,17 +36,14 @@
let
dm-haiku = buildPythonPackage rec {
pname = "dm-haiku";
version = "0.0.13";
version = "0.0.15";
pyproject = true;
# ImportError: `haiku.experimental.flax` features require `flax` to be installed.
disabled = pythonAtLeast "3.13";
src = fetchFromGitHub {
owner = "deepmind";
repo = "dm-haiku";
tag = "v${version}";
hash = "sha256-RJpQ9BzlbQ4X31XoJFnsZASiaC9fP2AdyuTAGINhMxs=";
hash = "sha256-phJ0f+effHQzuAVtPBR0bY3C0p//LBY7k1ci4mXBGfU=";
};
patches = [
@@ -58,17 +55,6 @@ let
})
];
# AttributeError: jax.core.Var was removed in JAX v0.6.0. Use jax.extend.core.Var instead, and
# see https://docs.jax.dev/en/latest/jax.extend.html for details.
# Already on master: https://github.com/google-deepmind/dm-haiku/commit/cfe8480d253a93100bf5e2d24c40435a95399c96
# TODO: remove at the next release
postPatch = ''
substituteInPlace haiku/_src/jaxpr_info.py \
--replace-fail "jax.core.JaxprEqn" "jax.extend.core.JaxprEqn" \
--replace-fail "jax.core.Var" "jax.extend.core.Var" \
--replace-fail "jax.core.Jaxpr" "jax.extend.core.Jaxpr"
'';
build-system = [ setuptools ];
dependencies = [
@@ -97,11 +83,12 @@ let
dm-env
dm-haiku
dm-tree
flax
jaxlib
optax
pytest-xdist
pytestCheckHook
rlax
# rlax (broken dependency tensorflow-probability)
tensorflow
];
@@ -120,11 +107,24 @@ let
# other tests to fail.
# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
"test_doctest_haiku.experimental"
# AssertionError: 1 != 0 : 1 doctests failed
"test_doctest_haiku"
# ValueError: pmap wrapped function must be passed at least one argument containing an array,
# got empty *args=() and **kwargs={}
"test_equivalent_when_passing_transformed_fn2"
# AssertionError: ValueError not raised
"test_passing_function_to_transform_pmap_transform"
"test_passing_function_to_transform_pmap_transform_with_state"
];
disabledTestPaths = [
# Those tests requires a more recent version of tensorflow. The current one (2.13) is not enough.
"haiku/_src/integration/jax2tf_test.py"
# Require rlax which is unavailable as its dependency tensorflow-probability is broken
"examples/impala/actor_test.py"
"examples/impala/learner_test.py"
"examples/impala_lite_test.py"
];
doCheck = false;
@@ -143,6 +143,7 @@ let
meta = {
description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet";
homepage = "https://github.com/deepmind/dm-haiku";
changelog = "https://github.com/google-deepmind/dm-haiku/releases/tag/${src.tag}";
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [ ndl ];
};
pkgs/development/python-modules/equinox/default.nix
+5 -1
View File
@@ -57,7 +57,11 @@ buildPythonPackage rec {
pytestCheckHook
];
disabledTests = lib.optionals stdenv.hostPlatform.isDarwin [
disabledTests = [
# Failed: DID NOT WARN. No warnings of type (<class 'Warning'>,) were emitted.
"test_jax_transform_warn"
]
++ lib.optionals stdenv.hostPlatform.isDarwin [
# SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
"test_filter"
];
pkgs/development/python-modules/evosax/default.nix
+5 -3
View File
@@ -7,14 +7,11 @@
setuptools,
# dependencies
chex,
dotmap,
flax,
jax,
jaxlib,
matplotlib,
numpy,
pyyaml,
# tests
# brax, (unpackaged)
@@ -62,6 +59,11 @@ buildPythonPackage rec {
# Requires unpackaged gymnax
"test_env_ffw_rollout"
# TypeError: ShapedArray.__init__() got an unexpected keyword argument 'named_shape'
"test_base_api"
"test_run"
"test_run_scan"
# Tries to download a data set from the internet
"test_brax_problem_eval"
"test_brax_problem_init"
pkgs/development/python-modules/flax/default.nix
+12
View File
@@ -2,6 +2,7 @@
lib,
buildPythonPackage,
fetchFromGitHub,
fetchpatch,
# build-system
setuptools,
@@ -48,6 +49,17 @@ buildPythonPackage rec {
hash = "sha256-ioMj8+TuOFX3t9p3oVaywaOQPFBgvNcy7b/2WX/yvXA=";
};
patches = [
# Fixes TypeError: shard_map() got an unexpected keyword argument 'auto'
# TODO: remove when updating to the next release
# https://github.com/google/flax/pull/5020
(fetchpatch {
name = "jax-0.8.0-compat";
url = "https://github.com/google/flax/commit/5bf9b35ff03130f440a93a812fd8b47ec6a49add.patch";
hash = "sha256-KYpa1wDQMt77XIDGQEg/VuU/OPPNp2enGSA986TZSLQ=";
})
];
build-system = [
setuptools
setuptools-scm
pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
+2 -2
View File
@@ -50,8 +50,8 @@ buildPythonPackage rec {
.${stdenv.hostPlatform.system};
hash =
{
x86_64-linux = "sha256-OXdyaiozKwvTSDG96ytWUzY0QvMBLCmW/IgICq9rO60=";
aarch64-linux = "sha256-2H1mbQxSP62q23GU58J03MWg5/j40dfig1NT7zK+8Bw=";
x86_64-linux = "sha256-npnD5Ub2vgD/xqTt0s9bdtTFmqNKjYOW2COauP2Ov9M=";
aarch64-linux = "sha256-UAu8W9mrvcyyOP4TgbkAzVdyZFANFHA1H2jl29tSNcI=";
}
.${stdenv.hostPlatform.system};
};
pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+6 -6
View File
@@ -39,32 +39,32 @@ let
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp311";
hash = "sha256-rckk68ekXI00AOoBGNxwpwgrKobjVxFzjUA904FdCb8=";
hash = "sha256-cL6wFZM7RyEDSBonUh1NtBz1kgCHBfRBLL86HsrzXC4=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp311";
hash = "sha256-KnJ6ia5prCHB9Qk9jVrviaDmkuZrA0/JNMiszHLkApA=";
hash = "sha256-x+pdt1IZS4o/v0xWS7e7wk8k4HEJu1nmWf1LETFLunY=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp312";
hash = "sha256-goTnz39USQZgTxEXAqbwARqW338BE4eLOBvsCQUXJTY=";
hash = "sha256-VVbj9Iy+HfYI2dDTQAMIZJ3eyqjQ4kOUkMb4i0GAMcs=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-mKl1ZVOChY2HTWRxzpcZQxBgnQoqfEKDxuB+N5M7d2g=";
hash = "sha256-SdBMFULqpAi3utAXASC82oMqAvuUT8iT3OiCbublpJM=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-chLBLXW33FEnXycYJ99KbTeEMMBvZQ5sMcFi/pV5/xI=";
hash = "sha256-Rw6W5+pNsPN1jabn0VxF6BAmYpJP46hHIV7mXKySxo0=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-Xj4qpNch+wLdECgmKq6uwpWORbylxNNRKykVG1cMtCU=";
hash = "sha256-IGmVNdeaLM+X8SEGC2BAUzhvQRzLMMElxmtqWgJsiSE=";
};
};
in
pkgs/development/python-modules/jax/default.nix
+2 -2
View File
@@ -40,7 +40,7 @@ let
in
buildPythonPackage rec {
pname = "jax";
version = "0.7.2";
version = "0.8.0";
pyproject = true;
src = fetchFromGitHub {
@@ -48,7 +48,7 @@ buildPythonPackage rec {
repo = "jax";
# google/jax contains tags for jax and jaxlib. Only use jax tags!
tag = "jax-v${version}";
hash = "sha256-GBpHFjvF7SvxJafu7aVlTp0jxSo4jAi9oPeMg2B/P24=";
hash = "sha256-Ilcp4WW65SyqrqDGBRdnB25m7OCbrsfdtxWvl0uTjNw=";
};
build-system = [ setuptools ];
pkgs/development/python-modules/jaxlib/bin.nix
+10 -10
View File
@@ -18,7 +18,7 @@
}:
let
version = "0.7.2";
version = "0.8.0";
inherit (python) pythonVersion;
# As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the
@@ -49,49 +49,49 @@ let
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp311";
hash = "sha256-Q4IAYjXM7VnS95WsyYPBvtz7yk/qj5RhMR1hxqeTrmY=";
hash = "sha256-/1Pouvl49rfEB2IVr3jwupacrENO0vclZdh+OMI/AOc=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp311";
hash = "sha256-jKcAM1H76Mz6L6Wkk+wt+/LfkkQTBs9cO5cFCO7bkqs=";
hash = "sha256-Qa693vZ6VVpt4XQnpOZs5gpSioFYR+Ldltq85Xn3rPg=";
};
"3.11-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp311";
hash = "sha256-n7+Qr84w4066LqkppQb1kHvdQGI1gSLeSZzp5nGguh8=";
hash = "sha256-u2AqjCTGFMuMpu7tPnCnM9k5nGoviJAKAlJiPNZydrU=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp312";
hash = "sha256-EfMjGeZizP9mhZ6zk3VwUNiXG9iAvE3XDexkNNiQ+1k=";
hash = "sha256-LIZ1v4bjka/k+NhjCAvhoCTXNN/T3RN/eqjn8iCRrc0=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-m1oNNXSXYRoRPSB/ssGZfwGrehdYcHAIEiIPC8qjGCI=";
hash = "sha256-2D/4zxsHApljnNpPhCdwf2kFHchCHln7tzMFUjk3Vw0=";
};
"3.12-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp312";
hash = "sha256-vW0cU71HXg52ilSvmLFkL7SdcwTPBVzuux0B6J04ocs=";
hash = "sha256-9gqsD2Tp5wpc7zQf4pJoRRhpVRTHGtAANndLvtX3MS4=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-SdmWIEhu/9qHQAAkcjpFIwZWaZbj3nGe5jPwUiDR7nc=";
hash = "sha256-Zsb1dvVKY+0FL1xGm+9NtyP18FC4OewMQpVzARNBvVg=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-l8eT6Xvl3cc7PoXmzorTcJ6AVPdeohnMDLTwgFplrwY=";
hash = "sha256-s+rFA7kP/sxo8R+hIhM+7yxixTbbKOgB5DbX56m2e/g=";
};
"3.13-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp313";
hash = "sha256-4bPf6ZFYJfzgBuoJW4U/V2gYRcW/qAl13MN4iTY3H7A=";
hash = "sha256-X88zpWOfjxZKRzqceKH6Cy4VrD/L7NbZaqD4i/Jeprs=";
};
};
in
pkgs/development/python-modules/lerobot/default.nix
+1
View File
@@ -113,6 +113,7 @@ buildPythonPackage rec {
"test_backward_compatibility"
"test_dataset_initialization"
"test_factory"
"test_from_pretrained_nonexistent_path"
"test_policy_defaults"
"test_save_and_load_pretrained"