Unverified Commit 7b16d5d8 authored by Gaetan Lepage's avatar Gaetan Lepage Committed by Nick Cao
Browse files

python3Packages.jaxlib-bin: 0.4.4 -> 0.4.14

parent 22114b44

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

+51 −33
Original line number Diff line number Diff line
@@ -18,11 +18,12 @@ 
, autoPatchelfHook

, buildPythonPackage

, config

, cudnn ? cudaPackages.cudnn

, fetchPypi

, fetchurl

, flatbuffers

, isPy39

, jaxlib-build

, lib

, ml-dtypes

, python

, scipy

, stdenv
@@ -35,46 +36,57 @@ let 
  inherit (cudaPackages) cudatoolkit cudnn;

in



assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1";

assert cudaSupport -> lib.versionAtLeast cudnn.version "8.2";

assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;



let

  version = "0.4.4";

  version = "0.4.14";



  pythonVersion = python.pythonVersion;

  inherit (python) pythonVersion;



  # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.

  # When upgrading, you can get these hashes from prefetch.sh. See

  # https://github.com/google/jax/issues/12879 as to why this specific URL is

  # the correct index.

  cpuSrcs = {

    "x86_64-linux" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp310-cp310-manylinux2014_x86_64.whl";

      hash = "sha256-4VT909AB+ti5HzQvsaZWNY6MS/GItlVEFH9qeZnUuKQ=";

  # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the

  # official instructions recommend installing CPU-only versions via PyPI.

  cpuSrcs =

    let

      getSrcFromPypi = { platform, hash }: fetchPypi {

        inherit version platform hash;

        pname = "jaxlib";

        format = "wheel";

        # See the `disabled` attr comment below.

        dist = "cp310";

        python = "cp310";

        abi = "cp310";

      };

    "aarch64-darwin" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/mac/jaxlib-${version}-cp310-cp310-macosx_11_0_arm64.whl";

      hash = "sha256-wuOmoCeTldslSa0MommQeTe+RYKhUMam1ZXrgSov+8U=";

    in

    {

      "x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        hash = "sha256-nyylSZfqHeftlvVgJZFCN1ldjluZVJIYu4ZSsVxvXf8=";

      };

    "x86_64-darwin" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/mac/jaxlib-${version}-cp310-cp310-macosx_10_14_x86_64.whl";

      hash = "sha256-arfiTw8yafJwjRwJhKby2O7y3+4ksh3PjaKW9JgJ1ok=";

      "aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        hash = "sha256-La3wYbGCjWTl7krBD6BaBRqyBD8R530Lckbz0AWv0FM=";

      };

      "x86_64-darwin" = getSrcFromPypi {

        platform = "macosx_10_14_x86_64";

        hash = "sha256-hDg5+qisgtgOrdvbjxsUgI73cW6Aah8NLjhPe4kMAsM=";

      };

    };





  # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.

  # When upgrading, you can get these hashes from prefetch.sh. See

  # 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.cudnn82-cp310-cp310-manylinux2014_x86_64.whl";

    hash = "sha256-bJ62DdzuPSV311ZI2R/LJQ3fOkDibtz2+8wDKw31FLk=";

    url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl";

    hash = "sha256-CcQ5kjp4XfUX4/RwFY3T5G3kVKAeyoCTXu1Lo4O16Qo=";

  };



in

buildPythonPackage rec {

buildPythonPackage {

  pname = "jaxlib";

  inherit version;

  format = "wheel";



  # At the time of writing (2022-10-19), there are releases for <=3.10.

  # Supporting all of them is a pain, so we focus on 3.10, the current nixpkgs

  # python version.

  disabled = !(pythonVersion == "3.10");



  # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
@@ -87,9 +99,10 @@ buildPythonPackage rec { 


  # Prebuilt wheels are dynamically linked against things that nix can't find.

  # Run `autoPatchelfHook` to automagically fix them.

  nativeBuildInputs = lib.optionals cudaSupport [ autoPatchelfHook addOpenGLRunpath ];

  nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ]

    ++ lib.optionals cudaSupport [ addOpenGLRunpath ];

  # Dynamic link dependencies

  buildInputs = [ stdenv.cc.cc ];

  buildInputs = [ stdenv.cc.cc.lib ];



  # jaxlib contains shared libraries that open other shared libraries via dlopen

  # and these implicit dependencies are not recognized by ldd or
@@ -113,7 +126,12 @@ buildPythonPackage rec { 
    done

  '';



  propagatedBuildInputs = [ absl-py flatbuffers scipy ];

  propagatedBuildInputs = [

    absl-py

    flatbuffers

    ml-dtypes

    scipy

  ];



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

  # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for
@@ -123,7 +141,7 @@ buildPythonPackage rec { 
    ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas

  '';



  pythonImportsCheck = [ "jaxlib" ];

  inherit (jaxlib-build) pythonImportsCheck;



  meta = with lib; {

    description = "XLA library for JAX";

pkgs/development/python-modules/jaxlib/prefetch.sh

+15 −7
Original line number Diff line number Diff line 
version="$1"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp39-none-manylinux2010_x86_64.whl)"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp310-none-manylinux2010_x86_64.whl)"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp39-none-manylinux2010_x86_64.whl)"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp39-none-manylinux2010_x86_64.whl)"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp310-none-manylinux2010_x86_64.whl)"

nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp310-none-manylinux2010_x86_64.whl)"
 
#!/usr/bin/env bash



prefetch () {

    expr="(import <nixpkgs> { system = \"$1\"; config.cudaSupport = $2; }).python3.pkgs.jaxlib-bin.src.url"

    url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r)

    echo "$url"

    sha256=$(nix-prefetch-url "$url")

    nix hash to-sri --type sha256 "$sha256"

    echo

}



prefetch "x86_64-linux" "false"

prefetch "aarch64-darwin" "false"

prefetch "x86_64-darwin" "false"

prefetch "x86_64-linux" "true"