python312Packages.{jax,jaxlib}: 0.4.28 -> 0.4.38; use jaxlib-bin by default (#363130)
This commit is contained in:
Gaétan Lepage
@@ -258,6 +258,11 @@
|
||||
|
||||
- `siduck76-st` has been renamed to `st-snazzy`, like the project's [flake](https://github.com/siduck/st/blob/main/flake.nix).
|
||||
|
||||
- `python3Packages.jax` now directly depends on `python3Packages.jaxlib`.
|
||||
As a result, packages that depend on jax no longer need to include jaxlib to their dependencies.
|
||||
There is also a breaking change in the handling of CUDA. Instead of using a CUDA compatible jaxlib
|
||||
as before, you can use plugins like `python3Packages.jax-cuda12-plugin`.
|
||||
|
||||
|
||||
<!-- To avoid merge conflicts, consider adding your item at an arbitrary place in the list instead. -->
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
{
|
||||
lib,
|
||||
stdenv,
|
||||
buildPythonPackage,
|
||||
fetchurl,
|
||||
autoAddDriverRunpath,
|
||||
autoPatchelfHook,
|
||||
pypaInstallHook,
|
||||
wheelUnpackHook,
|
||||
cudaPackages,
|
||||
python,
|
||||
jaxlib,
|
||||
}:
|
||||
let
|
||||
inherit (jaxlib) version;
|
||||
inherit (cudaPackages) cudaVersion;
|
||||
|
||||
cudaLibPath = lib.makeLibraryPath (
|
||||
with cudaPackages;
|
||||
[
|
||||
(lib.getLib cuda_cudart) # libcudart.so
|
||||
(lib.getLib cuda_cupti) # libcupti.so
|
||||
(lib.getLib cudnn) # libcudnn.so
|
||||
(lib.getLib libcufft) # libcufft.so
|
||||
(lib.getLib libcusolver) # libcusolver.so
|
||||
(lib.getLib libcusparse) # libcusparse.so
|
||||
]
|
||||
);
|
||||
|
||||
# Find new releases at https://storage.googleapis.com/jax-releases
|
||||
# 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.
|
||||
|
||||
# upstream does not distribute jax-cuda12-pjrt 0.4.38 binaries for aarch64-linux
|
||||
srcs = {
|
||||
"x86_64-linux" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_x86_64.whl";
|
||||
hash = "sha256-g75MWfvPMAd6YAhdmOfVncc4sckeDWKOSsF3n94VrCs=";
|
||||
};
|
||||
"aarch64-linux" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_aarch64.whl";
|
||||
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
|
||||
};
|
||||
};
|
||||
in
|
||||
buildPythonPackage {
|
||||
pname = "jax-cuda12-pjrt";
|
||||
inherit version;
|
||||
pyproject = false;
|
||||
|
||||
src =
|
||||
srcs.${stdenv.hostPlatform.system}
|
||||
or (throw "jax-cuda12-pjrt: No src for ${stdenv.hostPlatform.system}");
|
||||
|
||||
nativeBuildInputs = [
|
||||
autoAddDriverRunpath
|
||||
autoPatchelfHook
|
||||
pypaInstallHook
|
||||
wheelUnpackHook
|
||||
];
|
||||
|
||||
# The following attributes (buildInputs, postInstall and preInstallCheck) are copied from jaxlib-0.4.28
|
||||
# but it does not recognize GPUs as of 2024-12-29
|
||||
# Dynamic link dependencies
|
||||
buildInputs = [ (lib.getLib stdenv.cc.cc) ];
|
||||
|
||||
# jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`.
|
||||
# Linking into $out is the least bad solution. See
|
||||
# * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
|
||||
# * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
|
||||
# for more info.
|
||||
postInstall = ''
|
||||
mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin
|
||||
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas
|
||||
'';
|
||||
|
||||
# jaxlib contains shared libraries that open other shared libraries via dlopen
|
||||
# and these implicit dependencies are not recognized by ldd or
|
||||
# autoPatchelfHook. That means we need to sneak them into rpath. This step
|
||||
# must be done after autoPatchelfHook and the automatic stripping of
|
||||
# artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
|
||||
# patchPhase.
|
||||
preInstallCheck = ''
|
||||
shopt -s globstar
|
||||
|
||||
for file in $out/**/*.so; do
|
||||
echo $file
|
||||
patchelf --add-rpath "${cudaLibPath}" "$file"
|
||||
done
|
||||
'';
|
||||
|
||||
# no tests
|
||||
doCheck = false;
|
||||
|
||||
pythonImportsCheck = [ "jax_plugins.xla_cuda12" ];
|
||||
|
||||
meta = {
|
||||
description = "JAX XLA PJRT Plugin for NVIDIA GPUs";
|
||||
homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda";
|
||||
sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
|
||||
license = lib.licenses.asl20;
|
||||
maintainers = with lib.maintainers; [ natsukium ];
|
||||
platforms = lib.attrNames srcs;
|
||||
# see CUDA compatibility matrix
|
||||
# https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
|
||||
broken =
|
||||
!(lib.versionAtLeast cudaVersion "12.1")
|
||||
|| !(lib.versionAtLeast cudaPackages.cudnn.version "9.1")
|
||||
|| true;
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
{
|
||||
lib,
|
||||
stdenv,
|
||||
buildPythonPackage,
|
||||
fetchPypi,
|
||||
autoAddDriverRunpath,
|
||||
autoPatchelfHook,
|
||||
pypaInstallHook,
|
||||
wheelUnpackHook,
|
||||
cudaPackages,
|
||||
python,
|
||||
jaxlib,
|
||||
jax-cuda12-pjrt,
|
||||
}:
|
||||
let
|
||||
inherit (cudaPackages) cudaVersion;
|
||||
inherit (jaxlib) version;
|
||||
|
||||
getSrcFromPypi =
|
||||
{
|
||||
platform,
|
||||
dist,
|
||||
hash,
|
||||
}:
|
||||
fetchPypi {
|
||||
inherit
|
||||
version
|
||||
platform
|
||||
dist
|
||||
hash
|
||||
;
|
||||
pname = "jax_cuda12_plugin";
|
||||
format = "wheel";
|
||||
python = dist;
|
||||
abi = dist;
|
||||
};
|
||||
|
||||
# upstream does not distribute jax-cuda12-plugin 0.4.38 binaries for aarch64-linux
|
||||
srcs = {
|
||||
"3.10-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-nULpmc1k3VZ8FJ7Wj3k5K6iGRDZCGLtjbNzvoBl8kv4=";
|
||||
};
|
||||
"3.10-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
|
||||
};
|
||||
"3.11-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-cEZUOG8OYAoCgdquqViCqmekfttoOTthsbFzx+jKdKg=";
|
||||
};
|
||||
"3.11-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
|
||||
};
|
||||
"3.12-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-Ufas/3Ew63LrsCU039NYGg9eoGlx3lLX68Ia1Nh/5x4=";
|
||||
};
|
||||
"3.12-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
|
||||
};
|
||||
"3.13-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-CSKKTCtEO3aozZqOwikGAInEzINuBiSWh1ptb9xm0x8=";
|
||||
};
|
||||
"3.13-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
|
||||
};
|
||||
};
|
||||
in
|
||||
buildPythonPackage {
|
||||
pname = "jax-cuda12-plugin";
|
||||
inherit version;
|
||||
pyproject = false;
|
||||
|
||||
src = (
|
||||
srcs."${python.pythonVersion}-${stdenv.hostPlatform.system}"
|
||||
or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}")
|
||||
);
|
||||
|
||||
nativeBuildInputs = [
|
||||
autoAddDriverRunpath
|
||||
autoPatchelfHook
|
||||
pypaInstallHook
|
||||
wheelUnpackHook
|
||||
];
|
||||
|
||||
dependencies = [ jax-cuda12-pjrt ];
|
||||
|
||||
pythonImportsCheck = [ "jax_cuda12_plugin" ];
|
||||
|
||||
meta = {
|
||||
description = "JAX Plugin for CUDA12";
|
||||
homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda";
|
||||
sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
|
||||
license = lib.licenses.asl20;
|
||||
maintainers = with lib.maintainers; [ natsukium ];
|
||||
platforms = lib.platforms.linux;
|
||||
# see CUDA compatibility matrix
|
||||
# https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
|
||||
broken =
|
||||
!(lib.versionAtLeast cudaVersion "12.1")
|
||||
|| !(lib.versionAtLeast cudaPackages.cudnn.version "9.1")
|
||||
|| true;
|
||||
};
|
||||
}
|
||||
@@ -1,25 +1,38 @@
|
||||
{
|
||||
lib,
|
||||
config,
|
||||
stdenv,
|
||||
blas,
|
||||
buildPythonPackage,
|
||||
callPackage,
|
||||
setuptools,
|
||||
importlib-metadata,
|
||||
fetchFromGitHub,
|
||||
jaxlib,
|
||||
jaxlib-bin,
|
||||
jaxlib-build,
|
||||
hypothesis,
|
||||
lapack,
|
||||
matplotlib,
|
||||
buildPythonPackage,
|
||||
fetchFromGitHub,
|
||||
cudaSupport ? config.cudaSupport,
|
||||
|
||||
# build-system
|
||||
setuptools,
|
||||
|
||||
# dependencies
|
||||
jaxlib,
|
||||
ml-dtypes,
|
||||
numpy,
|
||||
opt-einsum,
|
||||
scipy,
|
||||
|
||||
# optional-dependencies
|
||||
jax-cuda12-plugin,
|
||||
|
||||
# tests
|
||||
cloudpickle,
|
||||
hypothesis,
|
||||
matplotlib,
|
||||
pytestCheckHook,
|
||||
pytest-xdist,
|
||||
pythonOlder,
|
||||
scipy,
|
||||
stdenv,
|
||||
|
||||
# passthru
|
||||
callPackage,
|
||||
jax,
|
||||
jaxlib-build,
|
||||
jaxlib-bin,
|
||||
}:
|
||||
|
||||
let
|
||||
@@ -27,38 +40,41 @@ let
|
||||
in
|
||||
buildPythonPackage rec {
|
||||
pname = "jax";
|
||||
version = "0.4.28";
|
||||
version = "0.4.38";
|
||||
pyproject = true;
|
||||
|
||||
disabled = pythonOlder "3.9";
|
||||
|
||||
src = fetchFromGitHub {
|
||||
owner = "google";
|
||||
repo = "jax";
|
||||
# google/jax contains tags for jax and jaxlib. Only use jax tags!
|
||||
rev = "refs/tags/jax-v${version}";
|
||||
hash = "sha256-qSHPwi3is6Ts7pz5s4KzQHBMbcjGp+vAOsejW3o36Ek=";
|
||||
tag = "jax-v${version}";
|
||||
hash = "sha256-H8I9Mkz6Ia1RxJmnuJOSevLGHN2J8ey59ZTlFx8YfnA=";
|
||||
};
|
||||
|
||||
nativeBuildInputs = [ setuptools ];
|
||||
build-system = [ setuptools ];
|
||||
|
||||
# The version is automatically set to ".dev" if this variable is not set.
|
||||
# https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3
|
||||
JAX_RELEASE = "1";
|
||||
|
||||
# jaxlib is _not_ included in propagatedBuildInputs because there are
|
||||
# different versions of jaxlib depending on the desired target hardware. The
|
||||
# JAX project ships separate wheels for CPU, GPU, and TPU.
|
||||
propagatedBuildInputs = [
|
||||
dependencies = [
|
||||
jaxlib
|
||||
ml-dtypes
|
||||
numpy
|
||||
opt-einsum
|
||||
scipy
|
||||
] ++ lib.optional (pythonOlder "3.10") importlib-metadata;
|
||||
] ++ lib.optionals cudaSupport optional-dependencies.cuda;
|
||||
|
||||
optional-dependencies = rec {
|
||||
cuda = [ jax-cuda12-plugin ];
|
||||
cuda12 = cuda;
|
||||
cuda12_pip = cuda;
|
||||
cuda12_local = cuda;
|
||||
};
|
||||
|
||||
nativeCheckInputs = [
|
||||
cloudpickle
|
||||
hypothesis
|
||||
jaxlib
|
||||
matplotlib
|
||||
pytestCheckHook
|
||||
pytest-xdist
|
||||
@@ -71,11 +87,17 @@ buildPythonPackage rec {
|
||||
# which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
|
||||
# Not a big deal, this is how the JAX docs suggest running the test suite
|
||||
# anyhow.
|
||||
pytestFlagsArray = [
|
||||
"--numprocesses=4"
|
||||
"-W ignore::DeprecationWarning"
|
||||
"tests/"
|
||||
];
|
||||
pytestFlagsArray =
|
||||
[
|
||||
"--numprocesses=4"
|
||||
"-W ignore::DeprecationWarning"
|
||||
"tests/"
|
||||
]
|
||||
++ lib.optionals stdenv.hostPlatform.isDarwin [
|
||||
# SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
|
||||
"--deselect tests/shape_poly_test.py::ShapePolyTest"
|
||||
"--deselect tests/tree_util_test.py::TreeTest"
|
||||
];
|
||||
|
||||
# Prevents `tests/export_back_compat_test.py::CompatTest::test_*` tests from failing on darwin with
|
||||
# PermissionError: [Errno 13] Permission denied: '/tmp/back_compat_testdata/test_*.py'
|
||||
@@ -125,16 +147,27 @@ buildPythonPackage rec {
|
||||
# Fails on some hardware due to some numerical error
|
||||
# See https://github.com/google/jax/issues/18535
|
||||
"testQdwhWithOnRankDeficientInput5"
|
||||
]
|
||||
++ lib.optionals stdenv.hostPlatform.isDarwin [
|
||||
# SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
|
||||
"testInAxesPyTreePrefixMismatchError"
|
||||
"testInAxesPyTreePrefixMismatchErrorKwargs"
|
||||
"testOutAxesPyTreePrefixMismatchError"
|
||||
"test_tree_map"
|
||||
"test_vjp_rule_inconsistent_pytree_structures_error"
|
||||
"test_vmap_in_axes_tree_prefix_error"
|
||||
"test_vmap_mismatched_axis_sizes_error_message_issue_705"
|
||||
];
|
||||
|
||||
disabledTestPaths = [
|
||||
# Segmentation fault. See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba
|
||||
"tests/linalg_test.py"
|
||||
]
|
||||
++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
|
||||
# RuntimeWarning: invalid value encountered in cast
|
||||
"tests/lax_test.py"
|
||||
];
|
||||
disabledTestPaths =
|
||||
[
|
||||
# Segmentation fault. See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba
|
||||
"tests/linalg_test.py"
|
||||
]
|
||||
++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
|
||||
# RuntimeWarning: invalid value encountered in cast
|
||||
"tests/lax_test.py"
|
||||
];
|
||||
|
||||
pythonImportsCheck = [ "jax" ];
|
||||
|
||||
@@ -147,25 +180,26 @@ buildPythonPackage rec {
|
||||
#
|
||||
# NIXPKGS_ALLOW_UNFREE=1 nixglhost -- nix run --impure .#python3Packages.jax.passthru.tests.test_cuda_jaxlibBin
|
||||
passthru.tests = {
|
||||
test_cuda_jaxlibSource = callPackage ./test-cuda.nix {
|
||||
jaxlib = jaxlib-build.override { cudaSupport = true; };
|
||||
};
|
||||
# jaxlib-build is broken as of 2024-12-20
|
||||
# test_cuda_jaxlibSource = callPackage ./test-cuda.nix {
|
||||
# jax = jax.override { jaxlib = jaxlib-build; };
|
||||
# };
|
||||
test_cuda_jaxlibBin = callPackage ./test-cuda.nix {
|
||||
jaxlib = jaxlib-bin.override { cudaSupport = true; };
|
||||
jax = jax.override { jaxlib = jaxlib-bin; };
|
||||
};
|
||||
};
|
||||
|
||||
# updater fails to pick the correct branch
|
||||
passthru.skipBulkUpdate = true;
|
||||
|
||||
meta = with lib; {
|
||||
meta = {
|
||||
description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code";
|
||||
longDescription = ''
|
||||
This is the JAX frontend package, it's meant to be used together with one of the jaxlib implementations,
|
||||
e.g. `python3Packages.jaxlib`, `python3Packages.jaxlib-bin`, or `python3Packages.jaxlibWithCuda`.
|
||||
'';
|
||||
homepage = "https://github.com/google/jax";
|
||||
license = licenses.asl20;
|
||||
maintainers = with maintainers; [ samuela ];
|
||||
license = lib.licenses.asl20;
|
||||
maintainers = with lib.maintainers; [ samuela ];
|
||||
};
|
||||
}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
{
|
||||
jax,
|
||||
jaxlib,
|
||||
pkgs,
|
||||
}:
|
||||
|
||||
@@ -8,8 +7,7 @@ pkgs.writers.writePython3Bin "jax-test-cuda"
|
||||
{
|
||||
libraries = [
|
||||
jax
|
||||
jaxlib
|
||||
];
|
||||
] ++ jax.optional-dependencies.cuda;
|
||||
}
|
||||
''
|
||||
import jax
|
||||
|
||||
@@ -6,46 +6,24 @@
|
||||
|
||||
{
|
||||
absl-py,
|
||||
autoAddDriverRunpath,
|
||||
autoPatchelfHook,
|
||||
buildPythonPackage,
|
||||
config,
|
||||
fetchPypi,
|
||||
fetchurl,
|
||||
flatbuffers,
|
||||
jaxlib-build,
|
||||
lib,
|
||||
ml-dtypes,
|
||||
python,
|
||||
scipy,
|
||||
stdenv,
|
||||
# Options:
|
||||
cudaSupport ? config.cudaSupport,
|
||||
cudaPackages,
|
||||
}:
|
||||
|
||||
let
|
||||
inherit (cudaPackages) cudaVersion;
|
||||
|
||||
version = "0.4.28";
|
||||
|
||||
version = "0.4.38";
|
||||
inherit (python) pythonVersion;
|
||||
|
||||
cudaLibPath = lib.makeLibraryPath (
|
||||
with cudaPackages;
|
||||
[
|
||||
(lib.getLib cuda_cudart) # libcudart.so
|
||||
(lib.getLib cuda_cupti) # libcupti.so
|
||||
(lib.getLib cudnn) # libcudnn.so
|
||||
(lib.getLib libcufft) # libcufft.so
|
||||
(lib.getLib libcusolver) # libcusolver.so
|
||||
(lib.getLib libcusparse) # libcusparse.so
|
||||
]
|
||||
);
|
||||
|
||||
# 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 =
|
||||
srcs =
|
||||
let
|
||||
getSrcFromPypi =
|
||||
{
|
||||
@@ -68,182 +46,129 @@ let
|
||||
};
|
||||
in
|
||||
{
|
||||
"3.9-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp39";
|
||||
hash = "sha256-Slbr8FtKTBeRaZ2HTgcvP4CPCYa0AQsU+1SaackMqdw=";
|
||||
};
|
||||
"3.9-aarch64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_11_0_arm64";
|
||||
dist = "cp39";
|
||||
hash = "sha256-sBVi7IrXVxm30DiXUkiel+trTctMjBE75JFjTVKCrTw=";
|
||||
};
|
||||
"3.9-x86_64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_10_14_x86_64";
|
||||
dist = "cp39";
|
||||
hash = "sha256-T5jMg3srbG3P4Kt/+esQkxSSCUYRmqOvn6oTlxj/J4c=";
|
||||
};
|
||||
|
||||
"3.10-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-47zcb45g+FVPQVwU2TATTmAuPKM8OOVGJ0/VRfh1dps=";
|
||||
hash = "sha256-Ya7MuaJ8Z/24RQ9jVyQAGc1FEcudYqROR2R1bThIU60=";
|
||||
};
|
||||
"3.10-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-7hnBY6j98IOdTBi4il+/tOcxunxDdBbT5Ug+Vwu3ZOQ=";
|
||||
};
|
||||
"3.10-aarch64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_11_0_arm64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-8Djmi9ENGjVUcisLvjbmpEg4RDenWqnSg/aW8O2fjAk=";
|
||||
hash = "sha256-MLL1LLUNdHNK8vR3wlM6elg+O7eyyKzes2Hud9lAV3o=";
|
||||
};
|
||||
"3.10-x86_64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_10_14_x86_64";
|
||||
dist = "cp310";
|
||||
hash = "sha256-pCHSN/jCXShQFm0zRgPGc925tsJvUrxJZwS4eCKXvWY=";
|
||||
hash = "sha256-VcGbnT8zpvxZ9kSqWiH7oCY5zN13bLSptVJmJfV4Of8=";
|
||||
};
|
||||
|
||||
"3.11-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-Rc4PPIQM/4I2z/JsN/Jsn/B4aV+T4MFiwyDCgfUEEnU=";
|
||||
hash = "sha256-J1H/cDfWqZfQvg53zEvjgcWp+buLMU7bdVwTpv2Wn0U=";
|
||||
};
|
||||
"3.11-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-Q9tYxMQnYnKWNmpWwQMY4fAPUDaQ4X+Uu0NEKT4ZleA=";
|
||||
};
|
||||
"3.11-aarch64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_11_0_arm64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-eThX+vN/Nxyv51L+pfyBH0NeQ7j7S1AgWERKf17M+Ck=";
|
||||
hash = "sha256-P7DqrnNpFXr+y+rVCq8p5z/936d6IzXXIb2XlPPFEOQ=";
|
||||
};
|
||||
"3.11-x86_64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_10_14_x86_64";
|
||||
dist = "cp311";
|
||||
hash = "sha256-L/gpDtx7ksfq5SUX9lSSYz4mey6QZ7rT5MMj0hPnfPU=";
|
||||
hash = "sha256-tn/eq9bf7Qi3do873/tSEWAIX4MFZpvRl77vYdCN4Is=";
|
||||
};
|
||||
|
||||
"3.12-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-RqGqhX9P7uikP8upXA4Kti1AwmzJcwtsaWVZCLo1n40=";
|
||||
hash = "sha256-2tbAqWVnwG0IPARp/sQPIBIQsJk2W9aYvjGm0uyI/Vk=";
|
||||
};
|
||||
"3.12-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-SW9FsOABojQTCc0MdK8LZwU33O15wWjLIwz8x3PwqoY";
|
||||
};
|
||||
"3.12-aarch64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_11_0_arm64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-jdi//jhTcC9jzZJNoO4lc0pNGc1ckmvgM9dyun0cF10=";
|
||||
hash = "sha256-8zvK/jLJelYuz2iU18QWdMgMCs3t+lQj1Jr1EUcUmHQ=";
|
||||
};
|
||||
"3.12-x86_64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_10_14_x86_64";
|
||||
dist = "cp312";
|
||||
hash = "sha256-1sCaVFMpciRhrwVuc1FG0sjHTCKsdCaoRetp8ya096A=";
|
||||
hash = "sha256-P+/qmF8EFYFvO7r9PwOkNwUCde+brJpywTFOFkSsV8E=";
|
||||
};
|
||||
|
||||
"3.13-x86_64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_x86_64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-LOd7qM2pJZpLypevwcci5CkabEY6Y/jTcsbtyFEX1iU=";
|
||||
};
|
||||
"3.13-aarch64-linux" = getSrcFromPypi {
|
||||
platform = "manylinux2014_aarch64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-JIzKN3Hr8ksHD0lwE2TOraM+YTlEWwbHgsylrFrZK/Q=";
|
||||
};
|
||||
"3.13-aarch64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_11_0_arm64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-b+MmuK82Y4fdR8zzElg7Kxf+0ScSybdKZIsYoTy9ur8=";
|
||||
};
|
||||
"3.13-x86_64-darwin" = getSrcFromPypi {
|
||||
platform = "macosx_10_14_x86_64";
|
||||
dist = "cp313";
|
||||
hash = "sha256-QeVa5YGKiC5XiehI9vFmh6wTK8+7Wl+hFKXRi3jQXy0=";
|
||||
};
|
||||
};
|
||||
|
||||
# Note that the prebuilt jaxlib binary requires specific version of CUDA to
|
||||
# work. The cuda12 jaxlib binaries only works with CUDA 12.2, and cuda11
|
||||
# jaxlib binaries only works with CUDA 11.8. This is why we need to find a
|
||||
# binary that matches the provided cudaVersion.
|
||||
gpuSrcVersionString = "cuda${cudaVersion}-${pythonVersion}";
|
||||
|
||||
# Find new releases at https://storage.googleapis.com/jax-releases
|
||||
# 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.
|
||||
gpuSrcs = {
|
||||
"cuda12.2-3.9" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";
|
||||
hash = "sha256-d8LIl22gIvmWfoyKfXKElZJXicPQIZxdS4HumhwQGCw=";
|
||||
};
|
||||
"cuda12.2-3.10" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";
|
||||
hash = "sha256-PXtWv+UEcMWF8LhWe6Z1UGkf14PG3dkJ0Iop0LiimnQ=";
|
||||
};
|
||||
"cuda12.2-3.11" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";
|
||||
hash = "sha256-QO2WSOzmJ48VaCha596mELiOfPsAGLpGctmdzcCHE/o=";
|
||||
};
|
||||
"cuda12.2-3.12" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";
|
||||
hash = "sha256-ixWMaIChy4Ammsn23/3cCoala0lFibuUxyUr3tjfFKU=";
|
||||
};
|
||||
};
|
||||
in
|
||||
buildPythonPackage {
|
||||
pname = "jaxlib";
|
||||
inherit version;
|
||||
format = "wheel";
|
||||
|
||||
disabled =
|
||||
!(
|
||||
pythonVersion == "3.9"
|
||||
|| pythonVersion == "3.10"
|
||||
|| pythonVersion == "3.11"
|
||||
|| pythonVersion == "3.12"
|
||||
);
|
||||
|
||||
# See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
|
||||
src =
|
||||
if !cudaSupport then
|
||||
(cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}"
|
||||
or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
|
||||
)
|
||||
else
|
||||
gpuSrcs."${gpuSrcVersionString}" or (throw "jaxlib-bin: No gpuSrc for ${gpuSrcVersionString}");
|
||||
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 ]
|
||||
++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
|
||||
nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ];
|
||||
# Dynamic link dependencies
|
||||
buildInputs = [ (lib.getLib stdenv.cc.cc) ];
|
||||
|
||||
# jaxlib contains shared libraries that open other shared libraries via dlopen
|
||||
# and these implicit dependencies are not recognized by ldd or
|
||||
# autoPatchelfHook. That means we need to sneak them into rpath. This step
|
||||
# must be done after autoPatchelfHook and the automatic stripping of
|
||||
# artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
|
||||
# patchPhase.
|
||||
preInstallCheck = lib.optional cudaSupport ''
|
||||
shopt -s globstar
|
||||
|
||||
for file in $out/**/*.so; do
|
||||
patchelf --add-rpath "${cudaLibPath}" "$file"
|
||||
done
|
||||
'';
|
||||
|
||||
propagatedBuildInputs = [
|
||||
dependencies = [
|
||||
absl-py
|
||||
flatbuffers
|
||||
ml-dtypes
|
||||
scipy
|
||||
];
|
||||
|
||||
# jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`.
|
||||
# Linking into $out is the least bad solution. See
|
||||
# * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
|
||||
# * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
|
||||
# for more info.
|
||||
postInstall = lib.optional cudaSupport ''
|
||||
mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin
|
||||
ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas
|
||||
'';
|
||||
pythonImportsCheck = [ "jaxlib" ];
|
||||
|
||||
inherit (jaxlib-build) pythonImportsCheck;
|
||||
|
||||
meta = with lib; {
|
||||
meta = {
|
||||
description = "Prebuilt jaxlib backend from PyPi";
|
||||
homepage = "https://github.com/google/jax";
|
||||
sourceProvenance = with sourceTypes; [ binaryNativeCode ];
|
||||
license = licenses.asl20;
|
||||
maintainers = with maintainers; [ samuela ];
|
||||
platforms = [
|
||||
"aarch64-darwin"
|
||||
"x86_64-linux"
|
||||
"x86_64-darwin"
|
||||
];
|
||||
broken =
|
||||
!(cudaSupport -> lib.versionAtLeast cudaVersion "11.1")
|
||||
|| !(cudaSupport -> lib.versionAtLeast cudaPackages.cudnn.version "8.2")
|
||||
|| !(cudaSupport -> stdenv.hostPlatform.isLinux)
|
||||
|| !(cudaSupport -> (gpuSrcs ? "cuda${cudaVersion}-${pythonVersion}"))
|
||||
sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
|
||||
license = lib.licenses.asl20;
|
||||
maintainers = with lib.maintainers; [ samuela ];
|
||||
badPlatforms = [
|
||||
# Fails at pythonImportsCheckPhase:
|
||||
# ...-python-imports-check-hook.sh/nix-support/setup-hook: line 10: 28017 Illegal instruction: 4
|
||||
# /nix/store/5qpssbvkzfh73xih07xgmpkj5r565975-python3-3.11.9/bin/python3.11 -c
|
||||
# 'import os; import importlib; list(map(lambda mod: importlib.import_module(mod), os.environ["pythonImportsCheck"].split()))'
|
||||
|| (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64);
|
||||
"x86_64-darwin"
|
||||
];
|
||||
};
|
||||
}
|
||||
|
||||
@@ -2,17 +2,23 @@
|
||||
#! nix-shell -i sh -p jq
|
||||
|
||||
prefetch () {
|
||||
expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = $3; }).python$1.pkgs.jaxlib-bin.src.url"
|
||||
expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.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"
|
||||
nix hash convert --to sri --hash-algo sha256 "$sha256"
|
||||
echo
|
||||
}
|
||||
|
||||
for py in "39" "310" "311" "312"; do
|
||||
prefetch "$py" "x86_64-linux" "false"
|
||||
prefetch "$py" "aarch64-darwin" "false"
|
||||
prefetch "$py" "x86_64-darwin" "false"
|
||||
prefetch "$py" "x86_64-linux" "true"
|
||||
for py in "310" "311" "312" "313"; do
|
||||
prefetch "$py" "x86_64-linux" "jaxlib-bin"
|
||||
prefetch "$py" "aarch64-linux" "jaxlib-bin"
|
||||
prefetch "$py" "aarch64-darwin" "jaxlib-bin"
|
||||
prefetch "$py" "x86_64-darwin" "jaxlib-bin"
|
||||
prefetch "$py" "x86_64-linux" "jax-cuda12-plugin"
|
||||
prefetch "$py" "aarch64-linux" "jax-cuda12-plugin"
|
||||
done
|
||||
|
||||
for arch in "x86_64-linux" "aarch64-linux"; do
|
||||
prefetch "312" "$arch" "jax-cuda12-pjrt"
|
||||
done
|
||||
|
||||
@@ -1,12 +1,9 @@
|
||||
{
|
||||
config,
|
||||
lib,
|
||||
stdenv,
|
||||
buildPythonPackage,
|
||||
fetchFromGitLab,
|
||||
|
||||
cudaSupport ? config.cudaSupport,
|
||||
|
||||
# build-system
|
||||
setuptools,
|
||||
|
||||
@@ -48,7 +45,7 @@ buildPythonPackage rec {
|
||||
ducc0
|
||||
h5py
|
||||
jax
|
||||
(jaxlib.override { inherit cudaSupport; })
|
||||
jaxlib
|
||||
matplotlib
|
||||
mpi4py
|
||||
mpi
|
||||
|
||||
@@ -6560,11 +6560,13 @@ self: super: with self; {
|
||||
|
||||
jax = callPackage ../development/python-modules/jax { };
|
||||
|
||||
jax-cuda12-pjrt = callPackage ../development/python-modules/jax-cuda12-pjrt { };
|
||||
|
||||
jax-cuda12-plugin = callPackage ../development/python-modules/jax-cuda12-plugin { };
|
||||
|
||||
jax-jumpy = callPackage ../development/python-modules/jax-jumpy { };
|
||||
|
||||
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
|
||||
inherit (pkgs.config) cudaSupport;
|
||||
};
|
||||
jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { };
|
||||
|
||||
jaxlib-build = callPackage ../development/python-modules/jaxlib rec {
|
||||
# Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.
|
||||
@@ -6572,16 +6574,7 @@ self: super: with self; {
|
||||
IOKit = pkgs.darwin.apple_sdk_11_0.IOKit;
|
||||
};
|
||||
|
||||
# Use the -bin on macOS since the source build doesn't support it (see #323154)
|
||||
jaxlib = if jaxlib-build.meta.unsupported then jaxlib-bin else jaxlib-build;
|
||||
|
||||
jaxlibWithCuda = self.jaxlib.override {
|
||||
cudaSupport = true;
|
||||
};
|
||||
|
||||
jaxlibWithoutCuda = self.jaxlib.override {
|
||||
cudaSupport = false;
|
||||
};
|
||||
jaxlib = jaxlib-bin;
|
||||
|
||||
jaxopt = callPackage ../development/python-modules/jaxopt { };
|
||||
|
||||
|
||||
Reference in New Issue
Block a user