Gaetan Lepage
68bc3961e4
Diff: https://github.com/google/jax/compare/jax-v0.10.2...jax-v0.11.0 Changelog: https://docs.jax.dev/en/latest/changelog.html
134 lines
3.8 KiB
Nix
134 lines
3.8 KiB
Nix
# For the moment we only support the CPU and GPU backends of jaxlib. The TPU
|
|
# backend will require some additional work. Those wheels are located here:
|
|
# https://storage.googleapis.com/jax-releases/libtpu_releases.html.
|
|
|
|
# See `python3Packages.jax.passthru` for CUDA tests.
|
|
|
|
{
|
|
absl-py,
|
|
autoPatchelfHook,
|
|
buildPythonPackage,
|
|
fetchPypi,
|
|
flatbuffers,
|
|
lib,
|
|
ml-dtypes,
|
|
python,
|
|
scipy,
|
|
stdenv,
|
|
}:
|
|
|
|
let
|
|
version = "0.11.0";
|
|
inherit (python) pythonVersion;
|
|
|
|
# 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.
|
|
srcs =
|
|
let
|
|
getSrcFromPypi =
|
|
{
|
|
platform,
|
|
dist,
|
|
hash,
|
|
}:
|
|
fetchPypi {
|
|
inherit
|
|
version
|
|
platform
|
|
dist
|
|
hash
|
|
;
|
|
pname = "jaxlib";
|
|
format = "wheel";
|
|
# See the `disabled` attr comment below.
|
|
python = dist;
|
|
abi = dist;
|
|
};
|
|
in
|
|
{
|
|
"3.12-x86_64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_x86_64";
|
|
dist = "cp312";
|
|
hash = "sha256-fqncBtlPU23qvPMEUTFDvIn1EXPd/XhtRPisMD799kM=";
|
|
};
|
|
"3.12-aarch64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_aarch64";
|
|
dist = "cp312";
|
|
hash = "sha256-G3Z8TXrU38Y1gxO2OubYPDVx086WbeBNiq8p60RjN4Y=";
|
|
};
|
|
"3.12-aarch64-darwin" = getSrcFromPypi {
|
|
platform = "macosx_11_0_arm64";
|
|
dist = "cp312";
|
|
hash = "sha256-9b8LO/wu9HeFgARzidrLUa4ADlDGs7b5jRB4iScGbJc=";
|
|
};
|
|
|
|
"3.13-x86_64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_x86_64";
|
|
dist = "cp313";
|
|
hash = "sha256-M92UBcqwXuTz7MOpEokyPpvS8I0lmQ5dRf+4Uk9RlQY=";
|
|
};
|
|
"3.13-aarch64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_aarch64";
|
|
dist = "cp313";
|
|
hash = "sha256-ESoBxYRwen0OhjT9Vd1+//6BKWbmDC06yE6MJORXEzY=";
|
|
};
|
|
"3.13-aarch64-darwin" = getSrcFromPypi {
|
|
platform = "macosx_11_0_arm64";
|
|
dist = "cp313";
|
|
hash = "sha256-iLNJqpXkUsqjCmcjMg6T7HvRqySenwvykADaG1765zM=";
|
|
};
|
|
|
|
"3.14-x86_64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_x86_64";
|
|
dist = "cp314";
|
|
hash = "sha256-YSXahTJhBkG30+p5JiF82rUtzLKU72YVRV4rbsueLuY=";
|
|
};
|
|
"3.14-aarch64-linux" = getSrcFromPypi {
|
|
platform = "manylinux_2_27_aarch64";
|
|
dist = "cp314";
|
|
hash = "sha256-SDQpzH+n/Wogy1D2ZsjS1Jnvt+m6MWOBHAHwUa1gV8Q=";
|
|
};
|
|
"3.14-aarch64-darwin" = getSrcFromPypi {
|
|
platform = "macosx_11_0_arm64";
|
|
dist = "cp314";
|
|
hash = "sha256-fmwOI3SUO6cZH+GW4iBnt4mJTFjPbd0kyZca5kLeWO4=";
|
|
};
|
|
};
|
|
in
|
|
buildPythonPackage {
|
|
pname = "jaxlib";
|
|
inherit version;
|
|
format = "wheel";
|
|
|
|
__structuredAttrs = true;
|
|
|
|
# See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
|
|
src = (
|
|
srcs."${pythonVersion}-${stdenv.hostPlatform.system}"
|
|
or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
|
|
);
|
|
|
|
# Prebuilt wheels are dynamically linked against things that nix can't find.
|
|
# Run `autoPatchelfHook` to automagically fix them.
|
|
nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ];
|
|
# Dynamic link dependencies
|
|
buildInputs = [ (lib.getLib stdenv.cc.cc) ];
|
|
|
|
dependencies = [
|
|
absl-py
|
|
flatbuffers
|
|
ml-dtypes
|
|
scipy
|
|
];
|
|
|
|
pythonImportsCheck = [ "jaxlib" ];
|
|
|
|
meta = {
|
|
description = "Prebuilt jaxlib backend from PyPi";
|
|
homepage = "https://github.com/google/jax";
|
|
sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
|
|
license = lib.licenses.asl20;
|
|
maintainers = with lib.maintainers; [ samuela ];
|
|
};
|
|
}
|