pyrox.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 writableTmpDirAsHomeHook,
61 pythonAtLeast,
62}:
63
64buildPythonPackage rec {
65 pname = "torch-geometric";
66 version = "2.6.1";
67 pyproject = true;
68
69 src = fetchFromGitHub {
70 owner = "pyg-team";
71 repo = "pytorch_geometric";
72 tag = version;
73 hash = "sha256-Zw9YqPQw2N0ZKn5i5Kl4Cjk9JDTmvZmyO/VvIVr6fTU=";
74 };
75
76 build-system = [
77 flit-core
78 ];
79
80 dependencies = [
81 aiohttp
82 fsspec
83 jinja2
84 numpy
85 psutil
86 pyparsing
87 requests
88 torch
89 tqdm
90 ];
91
92 optional-dependencies = {
93 benchmark = [
94 matplotlib
95 networkx
96 pandas
97 protobuf
98 wandb
99 ];
100 dev = [
101 ipython
102 matplotlib-inline
103 pre-commit
104 torch-geometric
105 ];
106 full = [
107 ase
108 # captum
109 graphviz
110 h5py
111 matplotlib
112 networkx
113 numba
114 opt-einsum
115 pandas
116 pgmpy
117 pynndescent
118 # pytorch-memlab
119 rdflib
120 rdkit
121 scikit-image
122 scikit-learn
123 scipy
124 statsmodels
125 sympy
126 tabulate
127 torch-geometric
128 torchmetrics
129 trimesh
130 ];
131 graphgym = [
132 protobuf
133 pytorch-lightning
134 yacs
135 ];
136 modelhub = [
137 huggingface-hub
138 ];
139 test = [
140 onnx
141 onnxruntime
142 pytest
143 pytest-cov-stub
144 ];
145 };
146
147 pythonImportsCheck = [
148 "torch_geometric"
149 ];
150
151 nativeCheckInputs = [
152 pytestCheckHook
153 writableTmpDirAsHomeHook
154 ];
155
156 disabledTests = [
157 # RuntimeError: addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.
158 # PyTorch built with MKL has better support for addmm with sparse CPU tensors.
159 "test_asap"
160 "test_graph_unet"
161
162 # AttributeError: type object 'Any' has no attribute '_name'
163 "test_type_repr"
164
165 # AttributeError: module 'torch.fx._symbolic_trace' has no attribute 'List'
166 "test_set_clear_mask"
167 "test_sequential_to_hetero"
168 "test_to_fixed_size"
169 "test_to_hetero_basic"
170 "test_to_hetero_with_gcn"
171 "test_to_hetero_with_basic_model"
172 "test_to_hetero_and_rgcn_equal_output"
173 "test_graph_level_to_hetero"
174 "test_hetero_transformer_self_loop_error"
175 "test_to_hetero_validate"
176 "test_to_hetero_on_static_graphs"
177 "test_to_hetero_with_bases"
178 "test_to_hetero_with_bases_and_rgcn_equal_output"
179 "test_to_hetero_with_bases_validate"
180 "test_to_hetero_with_bases_on_static_graphs"
181 "test_to_hetero_with_bases_save"
182
183 # Failed: DID NOT WARN.
184 "test_to_hetero_validate"
185 "test_to_hetero_with_bases_validate"
186
187 # Failed: DID NOT RAISE
188 "test_scatter_backward"
189 ]
190 ++ lib.optionals stdenv.hostPlatform.isDarwin [
191 # This test uses `torch.jit` which might not be working on darwin:
192 # RuntimeError: required keyword attribute 'value' has the wrong type
193 "test_traceable_my_conv_with_self_loops"
194 ]
195 ++ lib.optionals (pythonAtLeast "3.13") [
196 # RuntimeError: Dynamo is not supported on Python 3.13+
197 "test_compile"
198
199 # RuntimeError: Python 3.13+ not yet supported for torch.compile
200 "test_compile_graph_breaks"
201 "test_compile_multi_aggr_sage_conv"
202 "test_compile_hetero_conv_graph_breaks"
203
204 # AttributeError: module 'typing' has no attribute 'io'. Did you mean: 'IO'?
205 "test_packaging"
206
207 # RuntimeError: Boolean value of Tensor with more than one value is ambiguous
208 "test_feature_store"
209 ];
210
211 meta = {
212 description = "Graph Neural Network Library for PyTorch";
213 homepage = "https://github.com/pyg-team/pytorch_geometric";
214 changelog = "https://github.com/pyg-team/pytorch_geometric/blob/${src.rev}/CHANGELOG.md";
215 license = lib.licenses.mit;
216 maintainers = with lib.maintainers; [ GaetanLepage ];
217 };
218}