Merge pull request #309 from zaz/support-cu128

Support CUDA 12.8

Author: gbg141

README.md

@@ -67,7 +67,7 @@ TopoBench now uses [**uv**](https://docs.astral.sh/uv/), an extremely fast Pytho
 3.  **Initialize Environment**:
     Use our centralized setup script to handle Python 3.11 virtualization and specialized hardware (CUDA) mapping.
     ```bash
-    # Usage: source uv_env_setup.sh [cpu|cu118|cu121]
+    # Usage: source uv_env_setup.sh [cpu|cu118|cu121|cu128]
     source uv_env_setup.sh cpu
     ```
     *This script performs the following:*

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 requires-python = ">=3.11, <3.12"
 
 dependencies=[
-    "torch==2.3.0",
+    "torch>=2.3.0",
     "tqdm",
     "charset-normalizer",
     "numpy",
@@ -56,12 +56,16 @@ dependencies=[
     "topomodelx @ git+https://github.com/pyt-team/TopoModelX.git",
     "toponetx @ git+https://github.com/pyt-team/TopoNetX.git@c378925",
     "lightning==2.4.0",
+]
+
+[project.optional-dependencies]
+# Required for NSD, ED-GNN, and point cloud lifting backbones.
+# Wheels must match your PyTorch + CUDA version; see find-links below.
+sparse = [
     "torch-scatter",
     "torch-sparse",
     "torch-cluster",
 ]
-
-[project.optional-dependencies]
 doc = [
     "jupyter",
     "nbsphinx",
@@ -87,7 +91,7 @@ test = [
 ]
 
 dev = ["TopoBench[test, lint]"]
-all = ["TopoBench[dev, doc]"]
+all = ["TopoBench[dev, doc, sparse]"]
 
 [project.urls]
 homepage="https://geometric-intelligence.github.io/topobench/index.html"
@@ -112,21 +116,22 @@ name = "pytorch-cu121"
 url = "https://download.pytorch.org/whl/cu121"
 explicit = true
 
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
 # Default find-links (will be overwritten by bash script)
 [tool.uv]
 find-links = ["https://data.pyg.org/whl/torch-2.3.0+cu121.html"]
+no-build-package = ["torch-scatter", "torch-sparse", "torch-cluster"]
 
 [tool.uv.sources]
 torch = [
     { index = "pytorch-cpu", marker = "sys_platform == 'darwin' or sys_platform == 'win32'" },
     { index = "pytorch-cu121", marker = "sys_platform == 'linux'" },
 ]
 
-[tool.uv.extra-build-dependencies]
-torch-cluster = ["torch==2.3.0"]
-torch-scatter = ["torch==2.3.0"]
-torch-sparse = ["torch==2.3.0"]
-
 # ==============================================================================
 #  TOOL CONFIGS
 # ==============================================================================
@@ -172,7 +177,7 @@ disable_error_code = ["import-untyped"]
 plugins = "numpy.typing.mypy_plugin"
 
 [[tool.mypy.overrides]]
-module = ["torch_cluster.*","networkx.*","scipy.spatial","scipy.sparse","toponetx.classes.simplicial_complex"]
+module = ["torch_cluster.*","torch_sparse.*","torch_scatter.*","networkx.*","scipy.spatial","scipy.sparse","toponetx.classes.simplicial_complex"]
 ignore_missing_imports = true
 
 [tool.pytest.ini_options]

test/test_import.py

@@ -0,0 +1,32 @@
+"""Tests for package import behavior."""
+
+import subprocess
+import sys
+
+
+def test_import_without_sparse():
+    """Verify topobench imports without torch_sparse/scatter/cluster.
+
+    Runs in a subprocess that blocks the sparse imports via sys.modules
+    to avoid modifying the test environment.
+    """
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys;"
+            "sys.modules['torch_sparse'] = None;"
+            "sys.modules['torch_scatter'] = None;"
+            "sys.modules['torch_cluster'] = None;"
+            "import topobench;"
+            "from topobench.nn.backbones.graph import BACKBONE_CLASSES;"
+            "from topobench.nn.backbones.hypergraph import BACKBONE_CLASSES;"
+            "print('ok')",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, (
+        f"import topobench failed without sparse packages:\n{result.stderr}"
+    )
+    assert "ok" in result.stdout

topobench/data/utils/io_utils.py

@@ -12,7 +12,6 @@
 import torch_geometric
 from toponetx.classes import SimplicialComplex
 from torch_geometric.data import Data
-from torch_sparse import coalesce
 
 from topobench.data.utils import get_complex_connectivity
 
@@ -370,6 +369,8 @@ def load_hypergraph_pickle_dataset(data_dir, data_name):
     torch_geometric.data.Data
         Hypergraph dataset.
     """
+    from torch_sparse import coalesce
+
     data_dir = osp.join(data_dir, data_name)
 
     # Load node features:
@@ -478,6 +479,8 @@ def load_hypergraph_content_dataset(data_dir, data_name):
     torch_geometric.data.Data
         Hypergraph dataset.
     """
+    from torch_sparse import coalesce
+
     # data_dir = osp.join(data_dir, data_name)
 
     p2idx_features_labels = osp.join(data_dir, f"{data_name}.content")

topobench/dataloader/utils.py

@@ -5,7 +5,6 @@
 
 import torch
 import torch_geometric
-from torch_sparse import SparseTensor
 
 
 class DomainData(torch_geometric.data.Data):
@@ -70,6 +69,8 @@ def to_data_list(batch):
     list
         List of data objects.
     """
+    from torch_sparse import SparseTensor
+
     for key, _ in batch:
         if batch[key].is_sparse:
             sparse_data = batch[key].coalesce()

topobench/nn/backbones/graph/nsd_utils/inductive_discrete_models.py

@@ -12,7 +12,6 @@
 
 import torch
 import torch.nn.functional as F
-import torch_sparse
 from torch import nn
 
 from .laplacian_builders import (
@@ -102,6 +101,8 @@ def forward(self, x, edge_index):
         torch.Tensor
             Output node features of shape [num_nodes, output_dim].
         """
+        import torch_sparse
+
         # Get actual number of nodes dynamically
         actual_num_nodes = x.size(0)
 
@@ -281,6 +282,8 @@ def forward(self, x, edge_index):
         torch.Tensor
             Output node features of shape [num_nodes, output_dim].
         """
+        import torch_sparse
+
         # Get actual number of nodes dynamically
         actual_num_nodes = x.size(0)
 
@@ -453,6 +456,8 @@ def forward(self, x, edge_index):
         torch.Tensor
             Output node features of shape [num_nodes, output_dim].
         """
+        import torch_sparse
+
         # Get actual number of nodes dynamically
         actual_num_nodes = x.size(0)
 

topobench/nn/backbones/graph/nsd_utils/laplacian_builders.py

@@ -14,7 +14,6 @@
 import torch
 from torch import nn
 from torch_geometric.utils import degree
-from torch_scatter import scatter_add
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 from .laplace import (
@@ -170,6 +169,8 @@ def forward(self, maps):
         saved_tril_maps : torch.Tensor
             Saved lower triangular restriction maps for analysis.
         """
+        from torch_scatter import scatter_add
+
         assert len(maps.size()) == 2
         assert maps.size(1) == self.d
         left_idx, right_idx = self.left_right_idx
@@ -352,6 +353,8 @@ def forward(self, maps):
         saved_tril_maps : torch.Tensor
             Saved lower triangular transport maps for analysis.
         """
+        from torch_scatter import scatter_add
+
         left_idx, right_idx = self.left_right_idx
         tril_row, tril_col = self.vertex_tril_idx
         tril_indices, diag_indices = self.tril_indices, self.diag_indices

topobench/nn/backbones/hypergraph/edgnn.py

@@ -5,7 +5,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torch_geometric
-import torch_scatter
 
 
 class EDGNN(nn.Module):
@@ -482,6 +481,8 @@ def forward(self, X, vertex, edges, X0):
         Tensor
             Output features.
         """
+        import torch_scatter
+
         N = X.shape[-2]
 
         Xve = self.W1(X)[..., vertex, :]  # [nnz, C]
@@ -562,6 +563,8 @@ def forward(self, X, vertex, edges, X0, beta=1.0):
         Tensor
             Output features.
         """
+        import torch_scatter
+
         N = X.shape[-2]
 
         Xve = X[..., vertex, :]  # [nnz, C]
@@ -666,6 +669,8 @@ def forward(self, X, vertex, edges, X0):
         Tensor
             Output features.
         """
+        import torch_scatter
+
         N = X.shape[-2]
 
         Xve = self.W1(X[..., vertex, :])  # [nnz, C]

topobench/transforms/liftings/pointcloud2hypergraph/pointnet_lifting.py

@@ -2,7 +2,6 @@
 
 import torch
 import torch_geometric
-from torch_cluster import fps, radius
 
 from topobench.transforms.liftings.pointcloud2hypergraph.base import (
     PointCloud2HypergraphLifting,
@@ -45,6 +44,7 @@ def lift_topology(self, data: torch_geometric.data.Data) -> dict:
         dict
             The lifted topology.
         """
+        from torch_cluster import fps, radius
 
         batch = (
             torch.zeros(data.num_nodes, dtype=torch.long)

topobench/transforms/liftings/pointcloud2hypergraph/voronoi_lifting.py

@@ -2,7 +2,6 @@
 
 import torch
 import torch_geometric
-from torch_cluster import fps, knn
 
 from topobench.transforms.liftings.pointcloud2hypergraph.base import (
     PointCloud2HypergraphLifting,
@@ -37,6 +36,7 @@ def lift_topology(self, data: torch_geometric.data.Data) -> dict:
         dict
             The lifted topology.
         """
+        from torch_cluster import fps, knn
 
         # Sample FPS induced Voronoi graph
         support_idcs = fps(data.x, ratio=self.support_ratio)