Merge pull request #197600 from samuela/samuela/chex

chex, dm-tree, flax upgrades
This commit is contained in:
Samuel Ainsworth 2022-10-25 11:33:00 -07:00 committed by GitHub
commit b2399a8190
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 28 additions and 18 deletions

View File

@ -12,14 +12,14 @@
buildPythonPackage rec {
pname = "chex";
version = "0.1.4";
version = "0.1.5";
format = "setuptools";
src = fetchFromGitHub {
owner = "deepmind";
repo = pname;
rev = "refs/tags/v${version}";
hash = "sha256-FiZElESyOVu1VJxUDNiN6HVADiaubDrGdQHYp2CN8f4=";
hash = "sha256-FYB0IhokM74HBY2wOJYE6xJrHxHHWhDSAZpWFs6HFu0=";
};
propagatedBuildInputs = [
@ -39,6 +39,11 @@ buildPythonPackage rec {
pytestCheckHook
];
disabledTests = [
# See https://github.com/deepmind/chex/issues/204.
"test_uninspected_checks"
];
meta = with lib; {
description = "Chex is a library of utilities for helping to write reliable JAX code.";
homepage = "https://github.com/deepmind/chex";

View File

@ -13,14 +13,13 @@
buildPythonPackage rec {
pname = "dm-tree";
# As of 2021-12-29, the latest stable version still builds with Bazel.
version = "unstable-2021-12-20";
version = "0.1.7";
src = fetchFromGitHub {
owner = "deepmind";
repo = "tree";
rev = "b452e5c2743e7489b4ba7f16ecd51c516d7cd8e3";
sha256 = "1r187xwpvnnj98lyasngcv3lbxz0ziihpl5dbnjbfbjr0kh6z0j9";
rev = version;
hash = "sha256-rg6dcGcbTGfK3h4WAyhwCjgM3o64Jj2SImxNsZXJHHM=";
};
patches = [

View File

@ -17,25 +17,15 @@
buildPythonPackage rec {
pname = "flax";
version = "0.6.0";
version = "0.6.1";
src = fetchFromGitHub {
owner = "google";
repo = pname;
rev = "refs/tags/v${version}";
sha256 = "sha256-egTYYFZxhE/Kk7jXRi1HmjCjyFia2LoRigH042isDu0=";
sha256 = "sha256-fZiODo+izOwGjCCTvi11GvUG/VQL1DV9bNXKjvIIw4A=";
};
patches = [
# Bump rich dependency, should be fixed in releases after 0.6.0
# https://github.com/google/flax/pull/2407
(fetchpatch {
url = "https://github.com/google/flax/commit/72189153f9779022b97858ae747c23fbaf571e3d.patch";
sha256 = "sha256-hKOn/M7qpBM6R1RIJpnXpRoZgIHqkwQZApN4L0fBzIE=";
name = "bump_rich_dependency.patch";
})
];
buildInputs = [ jaxlib ];
propagatedBuildInputs = [
@ -76,6 +66,22 @@ buildPythonPackage rec {
"examples/*"
];
disabledTests = [
# See https://github.com/google/flax/issues/2554.
"test_async_save_checkpoints"
"test_jax_array0"
"test_jax_array1"
"test_keep0"
"test_keep1"
"test_optimized_lstm_cell_matches_regular"
"test_overwrite_checkpoints"
"test_save_restore_checkpoints_target_empty"
"test_save_restore_checkpoints_target_none"
"test_save_restore_checkpoints_target_singular"
"test_save_restore_checkpoints_w_float_steps"
"test_save_restore_checkpoints"
];
meta = with lib; {
description = "Neural network library for JAX";
homepage = "https://github.com/google/flax";