Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137# 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.
{ autoPatchelfHook, buildPythonPackage, fetchPypi, lib, ml-dtypes, numpy, python, scipy, stdenv,}:
let version = "0.12.0.1"; 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"; python = dist; abi = dist; }; in { "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; hash = "sha256-eDGOmseckpcvgcdwkPVXvbSlpJghdRxu4/apxmqUiEI="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; hash = "sha256-QIEiAcO81MsQS4j3H1j+xMvQwyupDlySs/qPQIGfZ3Q="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; hash = "sha256-2yzIWE19DRZik+8kdbcmmIoQ9vm1qOO443EkFeVXA4A="; };
"3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; hash = "sha256-/cVpmbppC9kV9wKWEyaYSOqvhdIx7Iq8RmENO4QMGJw="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; hash = "sha256-yN4crTgnExFPZX9AY4PtxT7p/aRnnbTisJ0d/v4cTkU="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; hash = "sha256-i+FgsLhHvFcWjkMFcTdW9toNC7omcV5k+dL0bAp6a0I="; };
"3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; hash = "sha256-tgeadmnWJOvULRVxM5Vm5iz+ciX0taILjhZ3PqmoNXo="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; hash = "sha256-YJ3v+uo4jR7Z0cUQTY9gQXkfAHca3DBfcxa3QvUC/Ws="; }; "3.14-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp314"; hash = "sha256-k7jaGD8CNWA8RgtTDajNbrI+p/HGrcT4KkOmA1SX1GY="; }; };inbuildPythonPackage { 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 = [ ml-dtypes numpy scipy ];
pythonImportsCheck = [ "jaxlib" ];
meta = { description = "Prebuilt jaxlib backend from PyPi"; homepage = "https://github.com/jax-ml/jax"; sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ samuela ]; # Keep in sync with `srcs` above platforms = [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ]; };}