Merge pull request #263416 from samuela/samuela/jaxlib-bin (f1775e64) · Commits · nix / nixpkgs

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

+78 −15

Original line number	Diff line number	Diff line
		@@ -47,39 +47,102 @@ let
		# official instructions recommend installing CPU-only versions via PyPI.
		cpuSrcs =
		let
		getSrcFromPypi = { platform, hash }: fetchPypi {
		inherit version platform hash;
		getSrcFromPypi = { platform, dist, hash }: fetchPypi {
		inherit version platform dist hash;
		pname = "jaxlib";
		format = "wheel";
		# See the `disabled` attr comment below.
		dist = "cp310";
		python = "cp310";
		abi = "cp310";
		python = dist;
		abi = dist;
		};
		in
		{
		"x86_64-linux" = getSrcFromPypi {
		"3.9-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux2014_x86_64";
		dist = "cp39";
		hash = "sha256-8bTrWutuK0qVnbkcwMfgBf414YdaLc3GK5IsCm/JNPE=";
		};
		"3.9-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp39";
		hash = "sha256-Tmv2iOqlNbZqw/rYjef6GmM0N18EA5JTt6T3lQe+4Rs=";
		};
		"3.9-x86_64-darwin" = getSrcFromPypi {
		platform = "macosx_10_14_x86_64";
		dist = "cp39";
		hash = "sha256-mDT1INLqPdCkxtMMFR0qHLOIZdWEy8Iuzw1/vOoECsA=";
		};

		"3.10-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux2014_x86_64";
		dist = "cp310";
		hash = "sha256-ksnY+CPEstact5lKjbSg+ZSPJtSt0Y0NFWEFufBCByk=";
		};
		"aarch64-darwin" = getSrcFromPypi {
		"3.10-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp310";
		hash = "sha256-O7dHvdKLKfNELGfF4TKy7N5EX6Ca7Zu8OtLXWvFykR8=";
		};
		"x86_64-darwin" = getSrcFromPypi {
		"3.10-x86_64-darwin" = getSrcFromPypi {
		platform = "macosx_10_14_x86_64";
		dist = "cp310";
		hash = "sha256-gqKMUZSXrt8sQtTAoQbzAfCzO8gM9Y1/tZpuJVWyN0Y=";
		};

		"3.11-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux2014_x86_64";
		dist = "cp311";
		hash = "sha256-m+NDzwXMNboNjDl2nLY+vqAoN2dQJZVWb1UQDpqqDPw=";
		};
		"3.11-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp311";
		hash = "sha256-zCOAjaWWCQT9Jnm1jjc1Rh5gemqy7ACtTKLM0MqSJzM=";
		};
		"3.11-x86_64-darwin" = getSrcFromPypi {
		platform = "macosx_10_14_x86_64";
		dist = "cp311";
		hash = "sha256-gOLIxkk+2hew2GqWu1WgMVEx1YEutx7Zod7QbwsuUVQ=";
		};

		"3.12-x86_64-linux" = getSrcFromPypi {
		platform = "manylinux2014_x86_64";
		dist = "cp312";
		hash = "sha256-BZTmkgNuV4nWtfbY4t/19aP43szZQEdgpFXh5qwGRXk=";
		};
		"3.12-aarch64-darwin" = getSrcFromPypi {
		platform = "macosx_11_0_arm64";
		dist = "cp312";
		hash = "sha256-aAMTrLXU9EYwPv+kdeyI88/D7b4NANB39Fn8vuXUqFA=";
		};
		"3.12-x86_64-darwin" = getSrcFromPypi {
		platform = "macosx_10_14_x86_64";
		dist = "cp312";
		hash = "sha256-KHzlIfa9KtYcHX+i/F/SKaYTpD4/XjHVu5j3BdRTUmc=";
		};
		};

		# Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.
		# When upgrading, you can get these hashes from prefetch.sh. See
		# https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.
		gpuSrc = fetchurl {
		gpuSrcs = {
		"3.9" = fetchurl {
		url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";
		hash = "sha256-WB5Vbr/XeYKXCP/3DIXF20jR6/1xE3huX1h5ow8ETl0=";
		};
		"3.10" = fetchurl {
		url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";
		hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o=";
		};
		"3.11" = fetchurl {
		url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";
		hash = "sha256-Q8ZtF2GCrG30GFbCeCZTWPmW2TBybeXzh2u+NRiYpx4=";
		};
		"3.12" = fetchurl {
		url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";
		hash = "sha256-lphkSDOJ9SwbO0hp/xC1bYn5fWgth9A9Iwsc9zV0buI=";
		};
		};

		in
		buildPythonPackage {
		@@ -87,15 +150,15 @@ buildPythonPackage {
		inherit version;
		format = "wheel";

		disabled = !(pythonVersion == "3.10");
		disabled = !(pythonVersion == "3.9" \|\| pythonVersion == "3.10" \|\| pythonVersion == "3.11" \|\| pythonVersion == "3.12");

		# See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
		src =
		if !cudaSupport then
		(
		cpuSrcs."${stdenv.hostPlatform.system}"
		cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}"
		or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
		) else gpuSrc;
		) else gpuSrcs."${pythonVersion}";

		# Prebuilt wheels are dynamically linked against things that nix can't find.
		# Run `autoPatchelfHook` to automagically fix them.

pkgs/development/python-modules/jaxlib/prefetch.sh

+9 −6

Original line number	Diff line number	Diff line
		#!/usr/bin/env bash
		#! /usr/bin/env nix-shell
		#! nix-shell -i sh -p jq

		prefetch () {
		expr="(import <nixpkgs> { system = \"$1\"; config.cudaSupport = $2; }).python3.pkgs.jaxlib-bin.src.url"
		expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = $3; }).python$1.pkgs.jaxlib-bin.src.url"
		url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" \| jq -r)
		echo "$url"
		sha256=$(nix-prefetch-url "$url")
		@@ -9,7 +10,9 @@ prefetch () {
		echo
		}

		prefetch "x86_64-linux" "false"
		prefetch "aarch64-darwin" "false"
		prefetch "x86_64-darwin" "false"
		prefetch "x86_64-linux" "true"
		for py in "39" "310" "311" "312"; do
		prefetch "$py" "x86_64-linux" "false"
		prefetch "$py" "aarch64-darwin" "false"
		prefetch "$py" "x86_64-darwin" "false"
		prefetch "$py" "x86_64-linux" "true"
		done