Commit 33de75c3 authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python311Packages.jaxlib: 0.4.19 -> 0.4.20 

Changelog: https://jax.readthedocs.io/en/latest/changelog.html#jaxlib-0-4-20-nov-2-2023
parent 153ad318

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

+6 −13
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.19";

  version = "0.4.20";



  meta = with lib; {

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

    "astor_archive"

    "astunparse_archive"

    "boringssl"

    # Not packaged in nixpkgs

    # "com_github_googleapis_googleapis"

    # "com_github_googlecloudplatform_google_cloud_cpp"
@@ -151,7 +150,7 @@ let 
      repo = "jax";

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

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

      hash = "sha256-l5uLPqhg/hqtO9oJSaioow5cH/0jKHDVziGezkfnVcc=";

      hash = "sha256-WLYXUtchOaA6SGnKuVhN9CmV06xMCLQTEuEtL13ttZU=";

    };



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



      sha256 = (if cudaSupport then {

        x86_64-linux = "sha256-Hw4uFvltH7nlNN3qAEcQ+IR2FAOjRkvwyWA3rCPi7Vo=";

        x86_64-linux = "sha256-QczClHxHElLZCqIZlHc3z3DXJ7rZQJaMs2XIb+lxarI=";

      } else {

        x86_64-linux = "sha256-LEugnFwTV3EyeTZWgMvXzHbgeDPdmuT3daXCXJRMYVY=";

        aarch64-linux = "sha256-59rv/3RjD8pnveBDZ33xZoNQxLmnhMocsKMgVfYoO70=";

        x86_64-linux = "sha256-mqiJe4u0NYh1PKCbQfbo0U2e9/kYiBqj98d+BPHFSxQ=";

        aarch64-linux = "sha256-EuLqamVBJ+qoVMCFIYUT846AghltZolfLGdtO9UeXSM=";

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

    };
@@ -293,13 +292,7 @@ let 
          --replace "/usr/bin/install_name_tool" "${cctools}/bin/install_name_tool"

        substituteInPlace ../output/external/rules_cc/cc/private/toolchain/unix_cc_configure.bzl \

          --replace "/usr/bin/libtool" "${cctools}/bin/libtool"

      '' + (if stdenv.cc.isGNU then ''

        sed -i 's@-lprotobuf@-l:libprotobuf.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

        sed -i 's@-lprotoc@-l:libprotoc.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

      '' else if stdenv.cc.isClang then ''

        sed -i 's@-lprotobuf@${pkgs.protobuf}/lib/libprotobuf.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

        sed -i 's@-lprotoc@${pkgs.protobuf}/lib/libprotoc.a@' ../output/external/xla/third_party/systemlibs/protobuf.BUILD

      '' else throw "Unsupported stdenv.cc: ${stdenv.cc}");

      '';

    };



    inherit meta;