diff --git a/pkgs/development/python-modules/torch-cluster/default.nix b/pkgs/development/python-modules/torch-cluster/default.nix new file mode 100644 index 000000000000..8a8dd1e1d328 --- /dev/null +++ b/pkgs/development/python-modules/torch-cluster/default.nix @@ -0,0 +1,82 @@ +{ + lib, + stdenv, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + torch, + + # buildInputs + pybind11, + + # dependencies + scipy, + + # tests + pytestCheckHook, + + # passthru + nix-update-script, +}: + +buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { + pname = "torch-cluster"; + version = "2.6.3-unstable-2026-03-26"; + pyproject = true; + __structuredAttrs = true; + + # Last stable release is from 2023 + # Development is still active, but nothing was properly tagged on GitHub or Pypi + # See: https://github.com/rusty1s/pytorch_cluster/issues/270 + src = fetchFromGitHub { + owner = "rusty1s"; + repo = "pytorch_cluster"; + rev = "af7b9f0af6b74be1594eb3d0a1685470cbb21265"; + hash = "sha256-2SXkk7m+feqk7uDir3Ov31TujyIUrRSEwONaiaq3Vvs="; + }; + + build-system = [ + setuptools + torch + ]; + + buildInputs = [ + pybind11 + ]; + + dependencies = [ + scipy + ]; + + pythonImportsCheck = [ "torch_cluster" ]; + + nativeCheckInputs = [ + pytestCheckHook + ]; + + # Otherwise python imports torch_cluster from /build/source instead of $out/..., which fails when + # trying to load the inexistant .so artifacts. + preCheck = '' + rm -rf torch_cluster + ''; + + disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + # aarch64-linux fails cpuinfo test, because /sys/devices/system/cpu/ does not exist in the sandbox: + # RuntimeError: Failed to initialize cpuinfo! + "test_fps" + "test_nearest" + ]; + + passthru.updateScript = nix-update-script { + extraArgs = [ "--version=branch" ]; + }; + + meta = { + description = "PyTorch Extension Library of Optimized Graph Cluster Algorithms"; + homepage = "https://github.com/rusty1s/pytorch_cluster"; + 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 358f80d4d22b..15ff0914d91e 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19740,6 +19740,8 @@ self: super: with self; { torch-c-dlpack-ext = callPackage ../development/python-modules/torch-c-dlpack-ext { }; + torch-cluster = callPackage ../development/python-modules/torch-cluster { }; + torch-einops-utils = callPackage ../development/python-modules/torch-einops-utils { }; torch-geometric = callPackage ../development/python-modules/torch-geometric { };