Gaetan Lepage
2026-06-12 21:16:36 +00:00
parent 7085145e08
commit 16776b808c
@@ -2,7 +2,6 @@
lib,
buildPythonPackage,
fetchFromGitHub,
fetchpatch,
# build-system
flit-core,
@@ -23,29 +22,17 @@
buildPythonPackage (finalAttrs: {
pname = "chex";
version = "0.1.91";
version = "0.1.92";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "deepmind";
owner = "google-deepmind";
repo = "chex";
tag = "v${finalAttrs.version}";
hash = "sha256-lJ9+kvG7dRtfDVgvkcJ9/jtnX0lMfxY4mmZ290y/74U=";
hash = "sha256-PM76Q72Bgyms7dROJkmlpPuDvtqjHLPTDkUYqo08T74=";
};
patches = [
# jax.device_put_replicated is removed in jax 0.10.0
# This fix was merged upstream -> remove when updating to the next release
(fetchpatch {
url = "https://github.com/google-deepmind/chex/commit/5fbd2c9a9936799daf92354e0307b9e88b9cc163.patch";
excludes = [
"chex/_src/variants.py"
];
hash = "sha256-ZTimSq7/yt2UEiWmLcfFBadX8+VcaxuPhkQJEyiEZlE=";
})
];
build-system = [
flit-core
];
@@ -67,19 +54,9 @@ buildPythonPackage (finalAttrs: {
pytestCheckHook
];
disabledTests = [
# Jax 0.8.2 incompatibility (reported at https://github.com/google-deepmind/chex/issues/422)
# AssertionError: AssertionError not raised
"test_assert_tree_is_on_device"
# AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not match ...
"test_assert_tree_is_on_host"
# AssertionError: [Chex] Assertion assert_tree_is_sharded failed: ...
"test_assert_tree_is_sharded"
];
meta = {
description = "Library of utilities for helping to write reliable JAX code";
homepage = "https://github.com/deepmind/chex";
homepage = "https://github.com/google-deepmind/chex";
changelog = "https://github.com/google-deepmind/chex/releases/tag/${finalAttrs.src.tag}";
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [ ndl ];