diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix index 6e39b1e872ff..51745dfdc89a 100644 --- a/pkgs/development/python-modules/dm-haiku/default.nix +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -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 ]; }; diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix index d21a6774342c..d96939cb3a92 100644 --- a/pkgs/development/python-modules/equinox/default.nix +++ b/pkgs/development/python-modules/equinox/default.nix @@ -57,7 +57,11 @@ buildPythonPackage rec { pytestCheckHook ]; - disabledTests = lib.optionals stdenv.hostPlatform.isDarwin [ + disabledTests = [ + # Failed: DID NOT WARN. No warnings of type (,) 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" ]; diff --git a/pkgs/development/python-modules/evosax/default.nix b/pkgs/development/python-modules/evosax/default.nix index c11d37113080..e0de3f633427 100644 --- a/pkgs/development/python-modules/evosax/default.nix +++ b/pkgs/development/python-modules/evosax/default.nix @@ -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" diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix index bab2e95c9a71..5def21410d92 100644 --- a/pkgs/development/python-modules/flax/default.nix +++ b/pkgs/development/python-modules/flax/default.nix @@ -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 diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index f4baed3c75a7..9b5ce1996105 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -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}; }; diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index a63c9554ce79..705bbea7b075 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -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 diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index c742f797d909..dfc221a76e56 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -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 ]; diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 0dadc56757f7..56c5831cfe67 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -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 diff --git a/pkgs/development/python-modules/lerobot/default.nix b/pkgs/development/python-modules/lerobot/default.nix index 2dbf6c1196ce..c7622536b604 100644 --- a/pkgs/development/python-modules/lerobot/default.nix +++ b/pkgs/development/python-modules/lerobot/default.nix @@ -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"