Unverified Commit cababf47 authored by Nick Cao's avatar Nick Cao Committed by GitHub
Browse files

Merge pull request #258951 from GaetanLepage/jax 

python310Packages.{jax,jaxlib,jaxlib-bin}: 0.4.16 -> 0.4.17
parents aeb18cc5 f7b2405b

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

+3 −3
Original line number Diff line number Diff line
@@ -27,8 +27,8 @@ let 
in

buildPythonPackage rec {

  pname = "jax";

  version = "0.4.16";

  format = "pyproject";

  version = "0.4.17";

  pyproject = true;



  disabled = pythonOlder "3.9";
@@ -37,7 +37,7 @@ buildPythonPackage rec { 
    repo = pname;

    # google/jax contains tags for jax and jaxlib. Only use jax tags!

    rev = "refs/tags/${pname}-v${version}";

    hash = "sha256-q+8CXGxK8JX0bUMK4KJB3qV/EaLHg68D1B5UrtRz0Eg=";

    hash = "sha256-Lxi/lBBq7VlsT6CgnXPFcwbRU+T8630rBdm693E2jok=";

  };



  nativeBuildInputs = [

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

+5 −5
Original line number Diff line number Diff line
@@ -39,7 +39,7 @@ in 
assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;



let

  version = "0.4.16";

  version = "0.4.17";



  inherit (python) pythonVersion;
@@ -60,15 +60,15 @@ let 
    {

      "x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        hash = "sha256-4XyaDnKEMhAbfPEvN3RCDEjXTWbOL6tWrTlyYeiboVs=";

        hash = "sha256-Fg/OaLgqeabFImUujdmhCqycANFZnLfhZmca2QmqE54=";

      };

      "aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        hash = "sha256-IG2pCui/Yj+LDMbQwBVlu7yl2llqnaxMzz/MtBvBr6U=";

        hash = "sha256-OSx3n5AsQ+Ggr0kVna/++bWvlSq6ABRj+Yz5WlnvF/8=";

      };

      "x86_64-darwin" = getSrcFromPypi {

        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.

  gpuSrc = fetchurl {

    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

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

+5 −5
Original line number Diff line number Diff line
@@ -54,7 +54,7 @@ let 
  inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl;



  pname = "jaxlib";

  version = "0.4.16";

  version = "0.4.17";



  meta = with lib; {

    description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";
@@ -151,7 +151,7 @@ let 
      repo = "jax";

      # google/jax contains tags for jax and jaxlib. Only use jaxlib tags!

      rev = "refs/tags/${pname}-v${version}";

      hash = "sha256-q+8CXGxK8JX0bUMK4KJB3qV/EaLHg68D1B5UrtRz0Eg=";

      hash = "sha256-Lxi/lBBq7VlsT6CgnXPFcwbRU+T8630rBdm693E2jok=";

    };



    nativeBuildInputs = [
@@ -264,10 +264,10 @@ let 
      ];



      sha256 = (if cudaSupport then {

        x86_64-linux = "sha256-6HkrEWAPjGPj4zRxahl0FLiV7WZO/6zsdCX8STfV5EE=";

        x86_64-linux = "sha256-nRvvFAuP/9D8BWWVPjuZijVtk+F9IrBBHsNc5Daluy4=";

      } else {

        x86_64-linux = "sha256-MDnuJwJ/xKnC72Qub0ETYj5uQB2r8/AgGm10oqmzzcc=";

        aarch64-linux = "sha256-aVUm612VNEsjZLDrtiOPTqSk1t+AhmOx+pOG3bZdOAw=";

        x86_64-linux = "sha256-pPIJOELN62GqUuaKpcpaqHu7wbJHiZgtb2PVUPRr1Ek=";

        aarch64-linux = "sha256-Q0PYZkOkUYUHVtSHZDlWitslDZbjNq6yRZv/ZkhTmyc=";

      }).${stdenv.system} or (throw "jaxlib: unsupported system: ${stdenv.system}");

    };