Commit fe108f01 authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python3Packages.jaxlib: 0.4.4 -> 0.4.12

parent 263a816a

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

+32 −36
Original line number Diff line number Diff line
@@ -4,7 +4,7 @@ 


  # Build-time dependencies:

, addOpenGLRunpath

, bazel_5

, bazel_6

, binutils

, buildBazelPackage

, buildPythonPackage
@@ -26,6 +26,7 @@ 
  # Python dependencies:

, absl-py

, flatbuffers

, ml-dtypes

, numpy

, scipy

, six
@@ -35,7 +36,6 @@ 
, giflib

, grpc

, libjpeg_turbo

, protobuf

, python

, snappy

, zlib
@@ -52,7 +52,7 @@ let 
  inherit (cudaPackages) backendStdenv cudatoolkit cudaFlags cudnn nccl;



  pname = "jaxlib";

  version = "0.4.4";

  version = "0.4.12";



  meta = with lib; {

    description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";
@@ -137,14 +137,15 @@ let 
  bazel-build = buildBazelPackage rec {

    name = "bazel-build-${pname}-${version}";



    bazel = bazel_5;

    # See https://github.com/google/jax/blob/main/.bazelversion for the latest.

    bazel = bazel_6;



    src = fetchFromGitHub {

      owner = "google";

      repo = "jax";

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

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

      hash = "sha256-DP68UwS9bg243iWU4MLHN0pwl8LaOcW3Sle1ZjsLOHo=";

      hash = "sha256-2JwEpzB5RwmBjGktppKhCpiaBM0AR20wfsRoQ33lh8Y=";

    };



    nativeBuildInputs = [
@@ -168,7 +169,7 @@ let 
      numpy

      openssl

      pkgs.flatbuffers

      protobuf

      pkgs.protobuf

      pybind11

      scipy

      six
@@ -187,7 +188,8 @@ let 
      rm -f .bazelversion

    '';



    bazelTargets = [ "//build:build_wheel" ];

    bazelRunTarget = "//build:build_wheel";

    runTargetFlags = [ "--output_path=$out" "--cpu=${arch}" ];



    removeRulesCC = false;
@@ -206,7 +208,11 @@ let 
      build --action_env=PYENV_ROOT

      build --python_path="${python}/bin/python"

      build --distinct_host_configuration=false

      build --define PROTOBUF_INCLUDE_PATH="${protobuf}/include"

      build --define PROTOBUF_INCLUDE_PATH="${pkgs.protobuf}/include"

    '' + lib.optionalString (stdenv.targetPlatform.avxSupport && stdenv.targetPlatform.isUnix) ''

      build --config=avx_posix

    '' + lib.optionalString mklSupport ''

      build --config=mkl_open_source_only

    '' + lib.optionalString cudaSupport ''

      build --action_env CUDA_TOOLKIT_PATH="${cudatoolkit_joined}"

      build --action_env CUDNN_INSTALL_PATH="${cudnn}"
@@ -233,7 +239,7 @@ let 
    fetchAttrs = {

      TF_SYSTEM_LIBS = lib.concatStringsSep "," tf_system_libs;

      # we have to force @mkl_dnn_v1 since it's not needed on darwin

      bazelTargets = bazelTargets ++ [ "@mkl_dnn_v1//:mkl_dnn" ];

      bazelTargets = [ bazelRunTarget "@mkl_dnn_v1//:mkl_dnn" ];

      bazelFlags = bazelFlags ++ [

        "--config=avx_posix"

      ] ++ lib.optionals cudaSupport [
@@ -248,9 +254,9 @@ let 


      sha256 =

        if cudaSupport then

          "sha256-O6bM7Lc8eaFyO4Xzl5/hvBrbPioI+Yeqx9yNC97fvKk="

          "sha256-wpucplv03HQHZ2gWhVq4R798ouPH99T3X4hbu7IRxj4="

        else

          "sha256-gLMJfJSQIdGGY2Ivx4IgDWg0hc+mxzlqY11CUkSWcjI=";

          "sha256-v2tCFifMBJbqweZQ2rsw707Zxehu+B+YtxFk1iHdDgc=";

    };



    buildAttrs = {
@@ -260,25 +266,13 @@ let 
        "nsync" # fails to build on darwin

      ]);



      bazelFlags = bazelFlags ++ lib.optionals (stdenv.targetPlatform.isx86_64 && stdenv.targetPlatform.isUnix) [

        "--config=avx_posix"

      ] ++ lib.optionals cudaSupport [

        "--config=cuda"

      ] ++ lib.optionals mklSupport [

        "--config=mkl_open_source_only"

      ];

      # Note: we cannot do most of this patching at `patch` phase as the deps are not available yet.

      # 1) Fix pybind11 include paths.

      # 2) Link protobuf from nixpkgs (through TF_SYSTEM_LIBS when using gcc) to prevent crashes on

      # 1) Link protobuf from nixpkgs (through TF_SYSTEM_LIBS when using gcc) to prevent crashes on

      #    loading multiple extensions in the same python program due to duplicate protobuf DBs.

      # 3) Patch python path in the compiler driver.

      preBuild = ''

        for src in ./jaxlib/*.{cc,h} ./jaxlib/cuda/*.{cc,h}; do

          sed -i 's@include/pybind11@pybind11@g' $src

        done

      '' + lib.optionalString cudaSupport ''

      # 2) Patch python path in the compiler driver.

      preBuild = lib.optionalString cudaSupport ''

        export NIX_LDFLAGS+=" -L${backendStdenv.nixpkgsCompatibleLibstdcxx}/lib"

        patchShebangs ../output/external/org_tensorflow/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc.tpl

        patchShebangs ../output/external/xla/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc.tpl

      '' + lib.optionalString stdenv.isDarwin ''

        # Framework search paths aren't added by bintools hook

        # https://github.com/NixOS/nixpkgs/pull/41914
@@ -288,16 +282,12 @@ let 
        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/org_tensorflow/third_party/systemlibs/protobuf.BUILD

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

        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@${protobuf}/lib/libprotobuf.a@' ../output/external/org_tensorflow/third_party/systemlibs/protobuf.BUILD

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

        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}");



      installPhase = ''

        ./bazel-bin/build/build_wheel --output_path=$out --cpu=${arch}

      '';

    };



    inherit meta;
@@ -344,13 +334,19 @@ buildPythonPackage { 
    grpc

    jsoncpp

    libjpeg_turbo

    ml-dtypes

    numpy

    scipy

    six

    snappy

  ];



  pythonImportsCheck = [ "jaxlib" ];

  pythonImportsCheck = [

    "jaxlib"

    # `import jaxlib` loads surprisingly little. These imports are actually bugs that appeared in the 0.4.11 upgrade.

    "jaxlib.cpu_feature_guard"

    "jaxlib.xla_client"

  ];



  # Without it there are complaints about libcudart.so.11.0 not being found

  # because RPATH path entries added above are stripped.

pkgs/top-level/python-packages.nix

+0 −1
Original line number Diff line number Diff line
@@ -5264,7 +5264,6 @@ self: super: with self; { 
    # Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.

    cudaSupport = pkgs.config.cudaSupport or false;

    IOKit = pkgs.darwin.apple_sdk_11_0.IOKit;

    protobuf = pkgs.protobuf3_20; # jaxlib-build 0.3.15 won't build with protobuf 3.21

  };



  jaxlib = self.jaxlib-build;