tjh.dev
1{ stdenv, buildPythonPackage
2, fetchPypi, isPy3k, linuxPackages
3, fastrlock, numpy, six, wheel, pytest, mock, setuptools
4, cudatoolkit, cudnn, nccl
5}:
6
7buildPythonPackage rec {
8 pname = "cupy";
9 version = "7.0.0";
10
11 src = fetchPypi {
12 inherit pname version;
13 sha256 = "7230b25052b2b9074cc338b44338b4bd43575a48ad7b9c2ce7f36d7e0d117012";
14 };
15
16 checkInputs = [
17 pytest
18 mock
19 ];
20
21 preConfigure = ''
22 export CUDA_PATH=${cudatoolkit}
23 '';
24
25 propagatedBuildInputs = [
26 cudatoolkit
27 cudnn
28 linuxPackages.nvidia_x11
29 nccl
30 fastrlock
31 numpy
32 six
33 setuptools
34 wheel
35 ];
36
37 # In python3, test was failed...
38 doCheck = !isPy3k;
39
40 enableParallelBuilding = true;
41
42 meta = with stdenv.lib; {
43 description = "A NumPy-compatible matrix library accelerated by CUDA";
44 homepage = https://cupy.chainer.org/;
45 license = licenses.mit;
46 platforms = [ "x86_64-linux" ];
47 maintainers = with maintainers; [ hyphon81 ];
48 };
49}