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 [
158 # RuntimeError: addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.
159 # PyTorch built with MKL has better support for addmm with sparse CPU tensors.
160 "test_asap"
161 "test_graph_unet"
162
163 # AttributeError: type object 'Any' has no attribute '_name'
164 "test_type_repr"
165 ]
166 ++ lib.optionals stdenv.hostPlatform.isDarwin [
167 # This test uses `torch.jit` which might not be working on darwin:
168 # RuntimeError: required keyword attribute 'value' has the wrong type
169 "test_traceable_my_conv_with_self_loops"
170 ]
171 ++ lib.optionals (pythonAtLeast "3.13") [
172 # RuntimeError: Dynamo is not supported on Python 3.13+
173 "test_compile"
174
175 # RuntimeError: Python 3.13+ not yet supported for torch.compile
176 "test_compile_graph_breaks"
177 "test_compile_multi_aggr_sage_conv"
178 "test_compile_hetero_conv_graph_breaks"
179
180 # AttributeError: module 'typing' has no attribute 'io'. Did you mean: 'IO'?
181 "test_packaging"
182
183 # RuntimeError: Boolean value of Tensor with more than one value is ambiguous
184 "test_feature_store"
185 ];
186
187 meta = {
188 description = "Graph Neural Network Library for PyTorch";
189 homepage = "https://github.com/pyg-team/pytorch_geometric";
190 changelog = "https://github.com/pyg-team/pytorch_geometric/blob/${src.rev}/CHANGELOG.md";
191 license = lib.licenses.mit;
192 maintainers = with lib.maintainers; [ GaetanLepage ];
193 };
194}