pyrox.dev / nixpkgs
lol
fork atom

Merge pull request #262149 from GaetanLepage/jax

python310Packages.{jax,jaxlib,jaxlib-bin}: 0.4.18 -> 0.4.19

authored by Samuel Ainsworth and committed by GitHub c340267c 46da009c

+18 -18
unified split
+3 -3
pkgs/development/python-modules/chex/default.nix
··· 15 16 buildPythonPackage rec { 17 pname = "chex"; 18 - version = "0.1.83"; 19 format = "setuptools"; 20 21 disabled = pythonOlder "3.9"; 22 23 src = fetchFromGitHub { 24 owner = "deepmind"; 25 - repo = pname; 26 rev = "refs/tags/v${version}"; 27 - hash = "sha256-iEachJf5NjOnkMWdP0aVQHWNPgUUBkMnzHKq3GP7t4w="; 28 }; 29 30 propagatedBuildInputs = [
··· 15 16 buildPythonPackage rec { 17 pname = "chex"; 18 + version = "0.1.84"; 19 format = "setuptools"; 20 21 disabled = pythonOlder "3.9"; 22 23 src = fetchFromGitHub { 24 owner = "deepmind"; 25 + repo = "chex"; 26 rev = "refs/tags/v${version}"; 27 + hash = "sha256-LsUMvSMVGjqZuFDcb+/61RtFxweeG6bSFzmJUUMv6rA="; 28 }; 29 30 propagatedBuildInputs = [
+3 -3
pkgs/development/python-modules/jax/default.nix
··· 27 in 28 buildPythonPackage rec { 29 pname = "jax"; 30 - version = "0.4.18"; 31 pyproject = true; 32 33 disabled = pythonOlder "3.9"; 34 35 src = fetchFromGitHub { 36 owner = "google"; 37 - repo = pname; 38 # google/jax contains tags for jax and jaxlib. Only use jax tags! 39 rev = "refs/tags/${pname}-v${version}"; 40 - hash = "sha256-rDvWHa8jYCAA9iKbWaFUXdE/9L7AepFiNzmqOcc/090="; 41 }; 42 43 nativeBuildInputs = [
··· 27 in 28 buildPythonPackage rec { 29 pname = "jax"; 30 + version = "0.4.19"; 31 pyproject = true; 32 33 disabled = pythonOlder "3.9"; 34 35 src = fetchFromGitHub { 36 owner = "google"; 37 + repo = "jax"; 38 # google/jax contains tags for jax and jaxlib. Only use jax tags! 39 rev = "refs/tags/${pname}-v${version}"; 40 + hash = "sha256-l5uLPqhg/hqtO9oJSaioow5cH/0jKHDVziGezkfnVcc="; 41 }; 42 43 nativeBuildInputs = [
+5 -5
pkgs/development/python-modules/jaxlib/bin.nix
··· 39 assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux; 40 41 let 42 - version = "0.4.18"; 43 44 inherit (python) pythonVersion; 45 ··· 60 { 61 "x86_64-linux" = getSrcFromPypi { 62 platform = "manylinux2014_x86_64"; 63 - hash = "sha256-MpNomovvSVx4N6gsowOLksTyEgTK261vSXMGxYqlVOE="; 64 }; 65 "aarch64-darwin" = getSrcFromPypi { 66 platform = "macosx_11_0_arm64"; 67 - hash = "sha256-if/5O5DQVHFdsLw9O1creZBx5j8ftE7fsWMMX1NjHP0="; 68 }; 69 "x86_64-darwin" = getSrcFromPypi { 70 platform = "macosx_10_14_x86_64"; 71 - hash = "sha256-4NeHA/0SGdmHXyDGxpK7oJc7dE1meR4LPjzbIwxloqU="; 72 }; 73 }; 74 ··· 78 # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. 79 gpuSrc = fetchurl { 80 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; 81 - hash = "sha256-p6BNvhhRzVDQdpEoIRau5JovC+eDjlW3bXrahtsGvmI="; 82 }; 83 84 in
··· 39 assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux; 40 41 let 42 + version = "0.4.19"; 43 44 inherit (python) pythonVersion; 45 ··· 60 { 61 "x86_64-linux" = getSrcFromPypi { 62 platform = "manylinux2014_x86_64"; 63 + hash = "sha256-ksnY+CPEstact5lKjbSg+ZSPJtSt0Y0NFWEFufBCByk="; 64 }; 65 "aarch64-darwin" = getSrcFromPypi { 66 platform = "macosx_11_0_arm64"; 67 + hash = "sha256-O7dHvdKLKfNELGfF4TKy7N5EX6Ca7Zu8OtLXWvFykR8="; 68 }; 69 "x86_64-darwin" = getSrcFromPypi { 70 platform = "macosx_10_14_x86_64"; 71 + hash = "sha256-gqKMUZSXrt8sQtTAoQbzAfCzO8gM9Y1/tZpuJVWyN0Y="; 72 }; 73 }; 74 ··· 78 # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. 79 gpuSrc = fetchurl { 80 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; 81 + hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o="; 82 }; 83 84 in
+5 -5
pkgs/development/python-modules/jaxlib/default.nix
··· 54 inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl; 55 56 pname = "jaxlib"; 57 - version = "0.4.18"; 58 59 meta = with lib; { 60 description = "JAX is Autograd and XLA, brought together for high-performance machine learning research."; ··· 151 repo = "jax"; 152 # google/jax contains tags for jax and jaxlib. Only use jaxlib tags! 153 rev = "refs/tags/${pname}-v${version}"; 154 - hash = "sha256-rDvWHa8jYCAA9iKbWaFUXdE/9L7AepFiNzmqOcc/090="; 155 }; 156 157 nativeBuildInputs = [ ··· 264 ]; 265 266 sha256 = (if cudaSupport then { 267 - x86_64-linux = "sha256-0CfGWlwKsUFP1DHUN6+6wX3cHr5x3TE6NbqYlV5me1E="; 268 } else { 269 - x86_64-linux = "sha256-sljmyIligXC7d9fdlpqR32xyMR0UslWs04gXJBD8FTA="; 270 - aarch64-linux = "sha256-eJ4KIkHdcA2EVvyBoNum2cOPcHPFoBOtUTAGufO8FJA="; 271 }).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}"); 272 }; 273
··· 54 inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl; 55 56 pname = "jaxlib"; 57 + version = "0.4.19"; 58 59 meta = with lib; { 60 description = "JAX is Autograd and XLA, brought together for high-performance machine learning research."; ··· 151 repo = "jax"; 152 # google/jax contains tags for jax and jaxlib. Only use jaxlib tags! 153 rev = "refs/tags/${pname}-v${version}"; 154 + hash = "sha256-l5uLPqhg/hqtO9oJSaioow5cH/0jKHDVziGezkfnVcc="; 155 }; 156 157 nativeBuildInputs = [ ··· 264 ]; 265 266 sha256 = (if cudaSupport then { 267 + x86_64-linux = "sha256-Z5cSgdRxdKxidaz4b1RlUF4rVcQiUTmQ1OorlBWlpt0="; 268 } else { 269 + x86_64-linux = "sha256-sn7p8FFHWIVdBWnsLsVj5jLiSaTlRm7s/qj2RqvQ3jU="; 270 + aarch64-linux = "sha256-oAYF5AeuPHTlwtpDMs2+tAhRAJH0yeSVnB7Ni7wmzS8="; 271 }).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}"); 272 }; 273
+2 -2
pkgs/development/python-modules/pytensor/default.nix
··· 26 27 buildPythonPackage rec { 28 pname = "pytensor"; 29 - version = "2.17.2"; 30 pyproject = true; 31 32 disabled = pythonOlder "3.9"; ··· 35 owner = "pymc-devs"; 36 repo = "pytensor"; 37 rev = "refs/tags/rel-${version}"; 38 - hash = "sha256-u1CbOjU3rQ6G3SSwYR3UlebymkupGMJWID4RH4v9PIk="; 39 }; 40 41 postPatch = ''
··· 26 27 buildPythonPackage rec { 28 pname = "pytensor"; 29 + version = "2.17.3"; 30 pyproject = true; 31 32 disabled = pythonOlder "3.9"; ··· 35 owner = "pymc-devs"; 36 repo = "pytensor"; 37 rev = "refs/tags/rel-${version}"; 38 + hash = "sha256-FufPCFzSjG8BrHes7t3XsdovX9gqUBG0gMDGKvkRkSA="; 39 }; 40 41 postPatch = ''