Commit d9a6568d authored by Samuel Ainsworth's avatar Samuel Ainsworth
Browse files

python3Packages.jaxlib: 0.3.15 -> 0.3.22

parent 55b3f2ad

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

+8 −6
Original line number Diff line number Diff line
@@ -53,7 +53,7 @@ let 
  inherit (cudaPackages) cudatoolkit cudnn nccl;



  pname = "jaxlib";

  version = "0.3.15";

  version = "0.3.22";



  meta = with lib; {

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

      repo = "jax";

      rev = "${pname}-v${version}";

      sha256 = "sha256-pIl7zzl82w5HHnJadH2vtCT4mYFd5YmM9iHC2GoJD6s=";

      hash = "sha256-bnczJ8ma/UMKhA5MUQ6H4az+Tj+By14ZTG6lQQwptQs=";

    };



    nativeBuildInputs = [
@@ -235,11 +235,11 @@ let 
    fetchAttrs = {

      sha256 =

        if cudaSupport then

          "sha256-tdO4YjO985zbittb16RFWgxgUBrHYQfv5gRsA4IAkTk="

          "sha256-Z9GDWGv+1YFyJjudyshZfeRJsKShoA1kIbNR3h3GxPQ="

        else if stdenv.isDarwin then

          "sha256-+XYxfXBCASueqDGg0Zqcmpf7zmemYM6xCE+x0rl3j34="

          "sha256-i3wiJHD4+pgTvDMhnYiQo9pdxxKItgYnc4/4wGt2NXM="

        else

          "sha256-La1wC8X5aGK5mXvYy/kO8n4J+zaRZEc/DAX5zaH1D5A=";

          "sha256-liRxmjwm0OmVMfgoGXx+nGBdW2fzzP/d4zmK6A59HAM=";

    };



    buildAttrs = {
@@ -293,7 +293,9 @@ buildPythonPackage { 
  inherit meta pname version;

  format = "wheel";



  src = "${bazel-build}/jaxlib-${version}-cp${builtins.replaceStrings ["."] [""] python.pythonVersion}-none-${platformTag}.whl";

  src =

    let cp = "cp${builtins.replaceStrings ["."] [""] python.pythonVersion}";

    in "${bazel-build}/jaxlib-${version}-${cp}-${cp}-${platformTag}.whl";



  # Note that cudatoolkit is necessary since jaxlib looks for "ptxas" in $PATH.

  # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for