python3Packages.jax: 0.7.1 -> 0.7.2 (#444138) (4df66320) · Commits · nix / nixpkgs

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

+2 −22

Original line number	Diff line number	Diff line
		@@ -3,7 +3,6 @@
		stdenv,
		buildPythonPackage,
		fetchFromGitHub,
		fetchpatch2,

		# build-system
		hatchling,
		@@ -23,35 +22,16 @@

		buildPythonPackage rec {
		pname = "equinox";
		version = "0.13.0";
		version = "0.13.1";
		pyproject = true;

		src = fetchFromGitHub {
		owner = "patrick-kidger";
		repo = "equinox";
		tag = "v${version}";
		hash = "sha256-zXgAuFGWKHShKodi9swnWIry4VU9s4pBhBRoK5KzaL0=";
		hash = "sha256-txgL5a+kKT28gAS8HianBgnnR+J25R2wrpRr8HEWCXA=";
		};

		patches = [
		# The following two patches have been merged upstream and should be removed when updating to the next release
		# They fix the incompatibilities with jax>=0.7.0

		# https://github.com/patrick-kidger/equinox/pull/1086
		(fetchpatch2 {
		name = "remove-deprecated-batching-NotMapped";
		url = "https://github.com/patrick-kidger/equinox/commit/6a6a441ced2fe64191a087752f1c2e71a6ce39f1.patch";
		hash = "sha256-tzHFjMI3gAIh5MPkdbmzsky/oFjDEbOIkPGQMQ+gcQQ=";
		})

		# https://github.com/patrick-kidger/equinox/pull/1082
		(fetchpatch2 {
		name = "allow-creating-weak-references-to-flatten";
		url = "https://github.com/patrick-kidger/equinox/commit/62b3c94ad56bdb63524702b320e977d2d93dbe72.patch";
		hash = "sha256-c1FKCnC3/okuP2VJV4h7sPRYQeYJZSdzEG5ETL2M35k=";
		})
		];

		# Relax speed constraints on tests that can fail on busy builders
		postPatch = ''
		substituteInPlace tests/test_while_loop.py \

pkgs/development/python-modules/jax-cuda12-pjrt/default.nix

+3 −3

Original line number	Diff line number	Diff line
		@@ -45,13 +45,13 @@ buildPythonPackage rec {
		platform =
		{
		x86_64-linux = "manylinux_2_27_x86_64";
		aarch64-linux = "manylinux2014_aarch64";
		aarch64-linux = "manylinux_2_27_aarch64";
		}
		.${stdenv.hostPlatform.system};
		hash =
		{
		x86_64-linux = "sha256-quPn80WASloSKaxxSMvRNEUMgWGYu7/4WnYiGC7H9ko=";
		aarch64-linux = "sha256-f8QFBJYN5W4jL5MmMKJ1fs8/hzZlsTmDF9Jfa3RF1WA=";
		x86_64-linux = "sha256-OXdyaiozKwvTSDG96ytWUzY0QvMBLCmW/IgICq9rO60=";
		aarch64-linux = "sha256-2H1mbQxSP62q23GU58J03MWg5/j40dfig1NT7zK+8Bw=";
		}
		.${stdenv.hostPlatform.system};
		};

pkgs/development/python-modules/jax-cuda12-plugin/default.nix

+9 −9

Original line number	Diff line number	Diff line
		@@ -39,32 +39,32 @@ let
		"3.11-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux_2_27_x86_64";
		dist = "cp311";
		hash = "sha256-Yw9aGxuorpSsC0K8N1IeGXBcmlRURWV5+NKY5FC2/tw=";
		hash = "sha256-rckk68ekXI00AOoBGNxwpwgrKobjVxFzjUA904FdCb8=";
		};
		"3.11-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp311";
		hash = "sha256-iAQd0J4Bm4JnTZ9xZ7eWtIy1so3pVHLBosA2OxZtS/Y=";
		hash = "sha256-KnJ6ia5prCHB9Qk9jVrviaDmkuZrA0/JNMiszHLkApA=";
		};
		"3.12-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux_2_27_x86_64";
		dist = "cp312";
		hash = "sha256-LPPm/mNDtbV2TTWJPON16z5qhZBIty3I9wCv7CFai6Y=";
		hash = "sha256-goTnz39USQZgTxEXAqbwARqW338BE4eLOBvsCQUXJTY=";
		};
		"3.12-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp312";
		hash = "sha256-M9ZGYPYVg1pg4M+Z656KMuz/1dIGYTV/RZEqHYFkMPE=";
		hash = "sha256-mKl1ZVOChY2HTWRxzpcZQxBgnQoqfEKDxuB+N5M7d2g=";
		};
		"3.13-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux_2_27_x86_64";
		dist = "cp313";
		hash = "sha256-dDC8RnptxbvBhvgcq5L0GhJy9aodFkbBu/rqkleD2oU=";
		hash = "sha256-chLBLXW33FEnXycYJ99KbTeEMMBvZQ5sMcFi/pV5/xI=";
		};
		"3.13-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp313";
		hash = "sha256-w8IAegYZmwlYMZdsamZcEvylWs1+yEq0SI4a5Y7stJQ=";
		hash = "sha256-Xj4qpNch+wLdECgmKq6uwpWORbylxNNRKykVG1cMtCU=";
		};
		};
		in

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

+6 −3

Original line number	Diff line number	Diff line
		@@ -40,7 +40,7 @@ let
		in
		buildPythonPackage rec {
		pname = "jax";
		version = "0.7.1";
		version = "0.7.2";
		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-k3yuWs+SOxrDXysVWCFR3wliATtro+aDdDXfFVtOTBo=";
		hash = "sha256-GBpHFjvF7SvxJafu7aVlTp0jxSo4jAi9oPeMg2B/P24=";
		};

		build-system = [ setuptools ];
		@@ -179,6 +179,9 @@ buildPythonPackage rec {
		description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code";
		homepage = "https://github.com/google/jax";
		license = lib.licenses.asl20;
		maintainers = with lib.maintainers; [ samuela ];
		maintainers = with lib.maintainers; [
		GaetanLepage
		samuela
		];
		};
		}

pkgs/development/python-modules/jaxlib/bin.nix

+13 −13

Original line number	Diff line number	Diff line
		@@ -18,7 +18,7 @@
		}:

		let
		version = "0.7.1";
		version = "0.7.2";
		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-6vX2j1O/TcuTtlElOFR2Z2JViOTzzK7vBIeI/RjYwNU=";
		hash = "sha256-Q4IAYjXM7VnS95WsyYPBvtz7yk/qj5RhMR1hxqeTrmY=";
		};
		"3.11-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp311";
		hash = "sha256-n7GJw7OUcMQ5T/yxi3HkfP/Fv4Xo/LHjNpJoaww+BN0=";
		hash = "sha256-jKcAM1H76Mz6L6Wkk+wt+/LfkkQTBs9cO5cFCO7bkqs=";
		};
		"3.11-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp311";
		hash = "sha256-PzLD5MFntzJ8NC6C09+EB5cU6gtDcYvocdA5mZZws8k=";
		hash = "sha256-n7+Qr84w4066LqkppQb1kHvdQGI1gSLeSZzp5nGguh8=";
		};

		"3.12-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux_2_27_x86_64";
		dist = "cp312";
		hash = "sha256-dKvTE1eX+CRA3TcRo1y6FsQw0bumVHS4W7cOQXM6Uuk=";
		hash = "sha256-EfMjGeZizP9mhZ6zk3VwUNiXG9iAvE3XDexkNNiQ+1k=";
		};
		"3.12-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp312";
		hash = "sha256-8PH1KVa4wlGKsACk09jCG+d34dR/kmugNkDjkQYaQe4=";
		hash = "sha256-m1oNNXSXYRoRPSB/ssGZfwGrehdYcHAIEiIPC8qjGCI=";
		};
		"3.12-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp312";
		hash = "sha256-EnwHxydwPl1Z+E9lUWm+yEn0Qi5S+FRjSc7MMKihPh0=";
		hash = "sha256-vW0cU71HXg52ilSvmLFkL7SdcwTPBVzuux0B6J04ocs=";
		};

		"3.13-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux_2_27_x86_64";
		dist = "cp313";
		hash = "sha256-s1D1Gahu/1pLHuAUx/qjZYX0fz1jeH0fPpvf/pzEGmY=";
		hash = "sha256-SdmWIEhu/9qHQAAkcjpFIwZWaZbj3nGe5jPwUiDR7nc=";
		};
		"3.13-aarch64-linux" = getSrcFromPypi {
		platform = "manylinux2014_aarch64";
		platform = "manylinux_2_27_aarch64";
		dist = "cp313";
		hash = "sha256-zmobpgGXZIcMJ1B6yhjlJpmK061L6i3WHhnQSZw7iwQ=";
		hash = "sha256-l8eT6Xvl3cc7PoXmzorTcJ6AVPdeohnMDLTwgFplrwY=";
		};
		"3.13-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp313";
		hash = "sha256-Kic3nl7Sl2WYDvMtTXf1fdDh3ZZYA6xnRVQESsI80+w=";
		hash = "sha256-4bPf6ZFYJfzgBuoJW4U/V2gYRcW/qAl13MN4iTY3H7A=";
		};
		};
		in