tjh.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6
7 # build-system
8 flit-core,
9
10 # dependencies
11 aiohttp,
12 fsspec,
13 jinja2,
14 numpy,
15 psutil,
16 pyparsing,
17 requests,
18 torch,
19 tqdm,
20
21 # optional-dependencies
22 matplotlib,
23 networkx,
24 pandas,
25 protobuf,
26 wandb,
27 ipython,
28 matplotlib-inline,
29 pre-commit,
30 torch-geometric,
31 ase,
32 # captum,
33 graphviz,
34 h5py,
35 numba,
36 opt-einsum,
37 pgmpy,
38 pynndescent,
39 # pytorch-memlab,
40 rdflib,
41 rdkit,
42 scikit-image,
43 scikit-learn,
44 scipy,
45 statsmodels,
46 sympy,
47 tabulate,
48 torchmetrics,
49 trimesh,
50 pytorch-lightning,
51 yacs,
52 huggingface-hub,
53 onnx,
54 onnxruntime,
55 pytest,
56 pytest-cov-stub,
57
58 # tests
59 pytestCheckHook,
60}:
61
62buildPythonPackage rec {
63 pname = "torch-geometric";
64 version = "2.6.1";
65 pyproject = true;
66
67 src = fetchFromGitHub {
68 owner = "pyg-team";
69 repo = "pytorch_geometric";
70 rev = "refs/tags/${version}";
71 hash = "sha256-Zw9YqPQw2N0ZKn5i5Kl4Cjk9JDTmvZmyO/VvIVr6fTU=";
72 };
73
74 build-system = [
75 flit-core
76 ];
77
78 dependencies = [
79 aiohttp
80 fsspec
81 jinja2
82 numpy
83 psutil
84 pyparsing
85 requests
86 torch
87 tqdm
88 ];
89
90 optional-dependencies = {
91 benchmark = [
92 matplotlib
93 networkx
94 pandas
95 protobuf
96 wandb
97 ];
98 dev = [
99 ipython
100 matplotlib-inline
101 pre-commit
102 torch-geometric
103 ];
104 full = [
105 ase
106 # captum
107 graphviz
108 h5py
109 matplotlib
110 networkx
111 numba
112 opt-einsum
113 pandas
114 pgmpy
115 pynndescent
116 # pytorch-memlab
117 rdflib
118 rdkit
119 scikit-image
120 scikit-learn
121 scipy
122 statsmodels
123 sympy
124 tabulate
125 torch-geometric
126 torchmetrics
127 trimesh
128 ];
129 graphgym = [
130 protobuf
131 pytorch-lightning
132 yacs
133 ];
134 modelhub = [
135 huggingface-hub
136 ];
137 test = [
138 onnx
139 onnxruntime
140 pytest
141 pytest-cov-stub
142 ];
143 };
144
145 pythonImportsCheck = [
146 "torch_geometric"
147 ];
148
149 nativeCheckInputs = [
150 pytestCheckHook
151 ];
152
153 preCheck = ''
154 export HOME=$(mktemp -d)
155 '';
156
157 disabledTests =
158 [
159 # TODO: try to re-enable when triton will have been updated to 3.0
160 # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
161 # LoweringException: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler'
162 "test_compile_hetero_conv_graph_breaks"
163 "test_compile_multi_aggr_sage_conv"
164
165 # RuntimeError: addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.
166 # PyTorch built with MKL has better support for addmm with sparse CPU tensors.
167 "test_asap"
168 "test_graph_unet"
169
170 # AttributeError: type object 'Any' has no attribute '_name'
171 "test_type_repr"
172 ]
173 ++ lib.optionals stdenv.hostPlatform.isDarwin [
174 # This test uses `torch.jit` which might not be working on darwin:
175 # RuntimeError: required keyword attribute 'value' has the wrong type
176 "test_traceable_my_conv_with_self_loops"
177 ];
178
179 meta = {
180 description = "Graph Neural Network Library for PyTorch";
181 homepage = "https://github.com/pyg-team/pytorch_geometric";
182 changelog = "https://github.com/pyg-team/pytorch_geometric/blob/${src.rev}/CHANGELOG.md";
183 license = lib.licenses.mit;
184 maintainers = with lib.maintainers; [ GaetanLepage ];
185 };
186}