python3Packages.chex: fix jax 0.10.0 compatibility (25f652d9) · Commits · nix / nixpkgs

pkgs/development/python-modules/chex/default.nix

+13 −0

Original line number	Diff line number	Diff line
		@@ -2,6 +2,7 @@
		lib,
		buildPythonPackage,
		fetchFromGitHub,
		fetchpatch,

		# build-system
		flit-core,
		@@ -33,6 +34,18 @@ buildPythonPackage (finalAttrs: {
		hash = "sha256-lJ9+kvG7dRtfDVgvkcJ9/jtnX0lMfxY4mmZ290y/74U=";
		};

		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
		];