Merge pull request #258951 from GaetanLepage/jax
python310Packages.{jax,jaxlib,jaxlib-bin}: 0.4.16 -> 0.4.17
This commit is contained in:
commit
cababf4784
|
@ -27,8 +27,8 @@ let
|
||||||
in
|
in
|
||||||
buildPythonPackage rec {
|
buildPythonPackage rec {
|
||||||
pname = "jax";
|
pname = "jax";
|
||||||
version = "0.4.16";
|
version = "0.4.17";
|
||||||
format = "pyproject";
|
pyproject = true;
|
||||||
|
|
||||||
disabled = pythonOlder "3.9";
|
disabled = pythonOlder "3.9";
|
||||||
|
|
||||||
|
@ -37,7 +37,7 @@ buildPythonPackage rec {
|
||||||
repo = pname;
|
repo = pname;
|
||||||
# google/jax contains tags for jax and jaxlib. Only use jax tags!
|
# google/jax contains tags for jax and jaxlib. Only use jax tags!
|
||||||
rev = "refs/tags/${pname}-v${version}";
|
rev = "refs/tags/${pname}-v${version}";
|
||||||
hash = "sha256-q+8CXGxK8JX0bUMK4KJB3qV/EaLHg68D1B5UrtRz0Eg=";
|
hash = "sha256-Lxi/lBBq7VlsT6CgnXPFcwbRU+T8630rBdm693E2jok=";
|
||||||
};
|
};
|
||||||
|
|
||||||
nativeBuildInputs = [
|
nativeBuildInputs = [
|
||||||
|
|
|
@ -39,7 +39,7 @@ in
|
||||||
assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;
|
assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;
|
||||||
|
|
||||||
let
|
let
|
||||||
version = "0.4.16";
|
version = "0.4.17";
|
||||||
|
|
||||||
inherit (python) pythonVersion;
|
inherit (python) pythonVersion;
|
||||||
|
|
||||||
|
@ -60,15 +60,15 @@ let
|
||||||
{
|
{
|
||||||
"x86_64-linux" = getSrcFromPypi {
|
"x86_64-linux" = getSrcFromPypi {
|
||||||
platform = "manylinux2014_x86_64";
|
platform = "manylinux2014_x86_64";
|
||||||
hash = "sha256-4XyaDnKEMhAbfPEvN3RCDEjXTWbOL6tWrTlyYeiboVs=";
|
hash = "sha256-Fg/OaLgqeabFImUujdmhCqycANFZnLfhZmca2QmqE54=";
|
||||||
};
|
};
|
||||||
"aarch64-darwin" = getSrcFromPypi {
|
"aarch64-darwin" = getSrcFromPypi {
|
||||||
platform = "macosx_11_0_arm64";
|
platform = "macosx_11_0_arm64";
|
||||||
hash = "sha256-IG2pCui/Yj+LDMbQwBVlu7yl2llqnaxMzz/MtBvBr6U=";
|
hash = "sha256-OSx3n5AsQ+Ggr0kVna/++bWvlSq6ABRj+Yz5WlnvF/8=";
|
||||||
};
|
};
|
||||||
"x86_64-darwin" = getSrcFromPypi {
|
"x86_64-darwin" = getSrcFromPypi {
|
||||||
platform = "macosx_10_14_x86_64";
|
platform = "macosx_10_14_x86_64";
|
||||||
hash = "sha256-x5DqsmHqEb7Dl7dnxT5N0l30GKt5OPZpq3HGX9MFKmo=";
|
hash = "sha256-1L4axL8b4a4c2PX02kFKbQ3o3jbPLv/bV1jU1neJYHg=";
|
||||||
};
|
};
|
||||||
};
|
};
|
||||||
|
|
||||||
|
@ -78,7 +78,7 @@ let
|
||||||
# https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.
|
# https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.
|
||||||
gpuSrc = fetchurl {
|
gpuSrc = fetchurl {
|
||||||
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl";
|
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl";
|
||||||
hash = "sha256-eLOprP2kv6roodwRKZXVZFQCD1wC26TSTEDJBjMu/Uo=";
|
hash = "sha256-Ctdlr8mvlMcTnBSiyjEEvle5AGr+o1v6OI7XIqcTENM=";
|
||||||
};
|
};
|
||||||
|
|
||||||
in
|
in
|
||||||
|
|
|
@ -54,7 +54,7 @@ let
|
||||||
inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl;
|
inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl;
|
||||||
|
|
||||||
pname = "jaxlib";
|
pname = "jaxlib";
|
||||||
version = "0.4.16";
|
version = "0.4.17";
|
||||||
|
|
||||||
meta = with lib; {
|
meta = with lib; {
|
||||||
description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";
|
description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";
|
||||||
|
@ -151,7 +151,7 @@ let
|
||||||
repo = "jax";
|
repo = "jax";
|
||||||
# google/jax contains tags for jax and jaxlib. Only use jaxlib tags!
|
# google/jax contains tags for jax and jaxlib. Only use jaxlib tags!
|
||||||
rev = "refs/tags/${pname}-v${version}";
|
rev = "refs/tags/${pname}-v${version}";
|
||||||
hash = "sha256-q+8CXGxK8JX0bUMK4KJB3qV/EaLHg68D1B5UrtRz0Eg=";
|
hash = "sha256-Lxi/lBBq7VlsT6CgnXPFcwbRU+T8630rBdm693E2jok=";
|
||||||
};
|
};
|
||||||
|
|
||||||
nativeBuildInputs = [
|
nativeBuildInputs = [
|
||||||
|
@ -264,10 +264,10 @@ let
|
||||||
];
|
];
|
||||||
|
|
||||||
sha256 = (if cudaSupport then {
|
sha256 = (if cudaSupport then {
|
||||||
x86_64-linux = "sha256-6HkrEWAPjGPj4zRxahl0FLiV7WZO/6zsdCX8STfV5EE=";
|
x86_64-linux = "sha256-nRvvFAuP/9D8BWWVPjuZijVtk+F9IrBBHsNc5Daluy4=";
|
||||||
} else {
|
} else {
|
||||||
x86_64-linux = "sha256-MDnuJwJ/xKnC72Qub0ETYj5uQB2r8/AgGm10oqmzzcc=";
|
x86_64-linux = "sha256-pPIJOELN62GqUuaKpcpaqHu7wbJHiZgtb2PVUPRr1Ek=";
|
||||||
aarch64-linux = "sha256-aVUm612VNEsjZLDrtiOPTqSk1t+AhmOx+pOG3bZdOAw=";
|
aarch64-linux = "sha256-Q0PYZkOkUYUHVtSHZDlWitslDZbjNq6yRZv/ZkhTmyc=";
|
||||||
}).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}");
|
}).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}");
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue