Commit 043cbacb authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python310Packages.jaxlib: 0.4.14 -> 0.4.15

parent 3dd248db

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

+9 −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.14";

  version = "0.4.16";



  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-0KnILQkahSiA1uuyT+kgy1XaCcZ3cpx1q114e2pecvg=";

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

    };



    nativeBuildInputs = [
@@ -203,6 +203,10 @@ let 
    GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";

    GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";



    # The version is automatically set to ".dev" if this variable is not set.

    # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3

    JAXLIB_RELEASE = "1";



    preConfigure = ''

      # dummy ldconfig

      mkdir dummy-ldconfig
@@ -260,10 +264,10 @@ let 
      ];



      sha256 = (if cudaSupport then {

        x86_64-linux = "sha256-L+d4umcN8eZQJS7NtbyMhFbbGUVd0a73GOYbZx3bW9Q=";

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

      } else {

        x86_64-linux = "sha256-V1giQbu70RYjbNsqudibiCgvtFNRIJ8XG75QtIzjM4g=";

        aarch64-linux = "sha256-DRU4aT7kQffhsOllgHtSlIsYOeLF4Sy5o5RR1CaTle0=";

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

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

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

    };