diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix index ed93783889d6..bdd741606efe 100644 --- a/pkgs/development/python-modules/chex/default.nix +++ b/pkgs/development/python-modules/chex/default.nix @@ -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 ];