From 3ca2a8c214b019e1cde566de384cc3cf257cf3ee Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 2 Oct 2024 16:09:54 +0200
Subject: [PATCH] python312Packages.torch-geometric: init at 2.6.1

---
 .../torch-geometric/default.nix               | 179 ++++++++++++++++++
 pkgs/top-level/python-packages.nix            |   2 +
 2 files changed, 181 insertions(+)
 create mode 100644 pkgs/development/python-modules/torch-geometric/default.nix

diff --git a/pkgs/development/python-modules/torch-geometric/default.nix b/pkgs/development/python-modules/torch-geometric/default.nix
new file mode 100644
index 000000000000..ed3bbdc73569
--- /dev/null
+++ b/pkgs/development/python-modules/torch-geometric/default.nix
@@ -0,0 +1,179 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  flit-core,
+
+  # dependencies
+  aiohttp,
+  fsspec,
+  jinja2,
+  numpy,
+  psutil,
+  pyparsing,
+  requests,
+  torch,
+  tqdm,
+
+  # optional-dependencies
+  matplotlib,
+  networkx,
+  pandas,
+  protobuf,
+  wandb,
+  ipython,
+  matplotlib-inline,
+  pre-commit,
+  torch-geometric,
+  ase,
+  # captum,
+  graphviz,
+  h5py,
+  numba,
+  opt-einsum,
+  pgmpy,
+  pynndescent,
+  # pytorch-memlab,
+  rdflib,
+  rdkit,
+  scikit-image,
+  scikit-learn,
+  scipy,
+  statsmodels,
+  sympy,
+  tabulate,
+  torchmetrics,
+  trimesh,
+  pytorch-lightning,
+  yacs,
+  huggingface-hub,
+  onnx,
+  onnxruntime,
+  pytest,
+  pytest-cov,
+
+  # tests
+  pytestCheckHook,
+}:
+
+buildPythonPackage rec {
+  pname = "torch-geometric";
+  version = "2.6.1";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "pyg-team";
+    repo = "pytorch_geometric";
+    rev = "refs/tags/${version}";
+    hash = "sha256-Zw9YqPQw2N0ZKn5i5Kl4Cjk9JDTmvZmyO/VvIVr6fTU=";
+  };
+
+  build-system = [
+    flit-core
+  ];
+
+  dependencies = [
+    aiohttp
+    fsspec
+    jinja2
+    numpy
+    psutil
+    pyparsing
+    requests
+    torch
+    tqdm
+  ];
+
+  optional-dependencies = {
+    benchmark = [
+      matplotlib
+      networkx
+      pandas
+      protobuf
+      wandb
+    ];
+    dev = [
+      ipython
+      matplotlib-inline
+      pre-commit
+      torch-geometric
+    ];
+    full = [
+      ase
+      # captum
+      graphviz
+      h5py
+      matplotlib
+      networkx
+      numba
+      opt-einsum
+      pandas
+      pgmpy
+      pynndescent
+      # pytorch-memlab
+      rdflib
+      rdkit
+      scikit-image
+      scikit-learn
+      scipy
+      statsmodels
+      sympy
+      tabulate
+      torch-geometric
+      torchmetrics
+      trimesh
+    ];
+    graphgym = [
+      protobuf
+      pytorch-lightning
+      yacs
+    ];
+    modelhub = [
+      huggingface-hub
+    ];
+    test = [
+      onnx
+      onnxruntime
+      pytest
+      pytest-cov
+    ];
+  };
+
+  pythonImportsCheck = [
+    "torch_geometric"
+  ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+  ];
+
+  preCheck = ''
+    export HOME=$(mktemp -d)
+  '';
+
+  disabledTests = [
+    # TODO: try to re-enable when triton will have been updated to 3.0
+    # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
+    # LoweringException: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler'
+    "test_compile_hetero_conv_graph_breaks"
+    "test_compile_multi_aggr_sage_conv"
+
+    # RuntimeError: addmm: computation on CPU is not implemented for SparseCsr + SparseCsr @ SparseCsr without MKL.
+    # PyTorch built with MKL has better support for addmm with sparse CPU tensors.
+    "test_asap"
+    "test_graph_unet"
+
+    # AttributeError: type object 'Any' has no attribute '_name'
+    "test_type_repr"
+  ];
+
+  meta = {
+    description = "Graph Neural Network Library for PyTorch";
+    homepage = "https://github.com/pyg-team/pytorch_geometric";
+    changelog = "https://github.com/pyg-team/pytorch_geometric/blob/${src.rev}/CHANGELOG.md";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index c1800dfbba93..eb62367cb06e 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -15692,6 +15692,8 @@ self: super: with self; {
 
   torch-audiomentations = callPackage ../development/python-modules/torch-audiomentations { };
 
+  torch-geometric = callPackage ../development/python-modules/torch-geometric { };
+
   torch-pitch-shift = callPackage ../development/python-modules/torch-pitch-shift { };
 
   torch-bin = callPackage ../development/python-modules/torch/bin.nix {