forked from mirrors/nixpkgs
46 lines
1.4 KiB
Nix
46 lines
1.4 KiB
Nix
|
# For the moment we only support the CPU backend of jaxlib. GPU and TPU backends
|
||
|
|
# require some additional work. Their wheels are not located on PyPI.
|
||
|
|
# * CPU/GPU: https://storage.googleapis.com/jax-releases/jax_releases.html
|
||
|
|
# * TPU: https://storage.googleapis.com/jax-releases/libtpu_releases.html
|
||
|
|
|
||
|
|
{ autoPatchelfHook, buildPythonPackage, fetchPypi, isPy39, lib, stdenv
|
||
|
|
# propagatedBuildInputs
|
||
|
|
, absl-py, flatbuffers, scipy
|
||
|
|
}:
|
||
|
|
|
||
|
|
buildPythonPackage rec {
|
||
|
|
pname = "jaxlib";
|
||
|
|
version = "0.1.70";
|
||
|
|
format = "wheel";
|
||
|
|
|
||
|
|
# At the time of writing (8/19/21), there are releases for 3.7-3.9. Supporting
|
||
|
|
# all of them is a pain, so we focus on 3.9, the current nixpkgs python3
|
||
|
|
# version.
|
||
|
|
disabled = !isPy39;
|
||
|
|
|
||
|
|
src = fetchPypi {
|
||
|
|
inherit pname version format;
|
||
|
|
dist = "cp39";
|
||
|
|
python = "cp39";
|
||
|
|
platform = "manylinux2010_x86_64";
|
||
|
|
sha256 = "sha256-mytMTqoavpuRawj52MU5/iFj27SGlm8DaoQ5vd/3bss=";
|
||
|
|
};
|
||
|
|
|
||
|
|
# Prebuilt wheels are dynamically linked against things that nix can't find.
|
||
|
|
# Run `autoPatchelfHook` to automagically fix them.
|
||
|
|
nativeBuildInputs = [ autoPatchelfHook ];
|
||
|
|
# Dynamic link dependencies
|
||
|
|
buildInputs = [ stdenv.cc.cc ];
|
||
|
|
|
||
|
|
# pip dependencies
|
||
|
|
propagatedBuildInputs = [ absl-py flatbuffers scipy ];
|
||
|
|
|
||
|
|
meta = with lib; {
|
||
|
|
description = "XLA library for JAX";
|
||
|
|
homepage = "https://github.com/google/jax";
|
||
|
|
license = licenses.asl20;
|
||
|
|
maintainers = with maintainers; [ samuela ];
|
||
|
|
platforms = [ "x86_64-linux" ];
|
||
|
|
};
|
||
|
|
}
|