Add first model

hyperbench/data/dataset.py

@@ -6,11 +6,10 @@
 import requests
 
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 from torch import Tensor
 from torch.utils.data import Dataset as TorchDataset
-from hyperbench.types.hypergraph import HIFHypergraph
-from hyperbench.types.hdata import HData
+from hyperbench.types import HData, HIFHypergraph
 from hyperbench.utils.hif_utils import validate_hif_json
 
 
@@ -137,14 +136,14 @@ def __getitem__(self, index: int | List[int]) -> HData:
             sampled_edge_index, sampled_node_ids, sampled_edge_ids
         )
 
-        new_node_features = self.hdata.x[sampled_node_ids]
+        new_x = self.hdata.x[sampled_node_ids]
 
         new_edge_attr = None
         if self.hdata.edge_attr is not None and len(sampled_edge_ids) > 0:
             new_edge_attr = self.hdata.edge_attr[sampled_edge_ids]
 
         return HData(
-            x=new_node_features,
+            x=new_x,
             edge_index=new_edge_index,
             edge_attr=new_edge_attr,
             num_nodes=len(sampled_node_ids),
@@ -185,8 +184,9 @@ def process(self) -> HData:
                 ]
             )
         else:
-            # Fallback to zeros if no numeric attributes
-            x = torch.zeros((num_nodes, 1), dtype=torch.float)
+            # Fallback to ones if no node features, 1 is better as it can help during
+            # training (e.g., avoid zero multiplication), especially in first epochs
+            x = torch.ones((num_nodes, 1), dtype=torch.float)
 
         # remap node and edge IDs to 0-based contiguous IDs
         # Use dict comprehension for faster lookups
@@ -322,7 +322,7 @@ def __sample_edge_index(
         node_ids = edge_index[0]
         edge_ids = edge_index[1]
 
-        sampled_node_ids = torch.tensor(sampled_node_ids_list)
+        sampled_node_ids = torch.tensor(sampled_node_ids_list, device=node_ids.device)
 
         # Find incidences where the node is in our sampled node set
         # Example: edge_index[0] = [0, 0, 1, 2, 3, 4], sampled_node_ids = [0, 3]
@@ -405,9 +405,11 @@ def __to_0based_ids(
         Returns:
             Tensor of 0-based ids.
         """
-        id_to_0based_id = torch.zeros(n, dtype=torch.long)
+        device = original_ids.device
+
+        id_to_0based_id = torch.zeros(n, dtype=torch.long, device=device)
         n_ids_to_keep = len(ids_to_keep)
-        id_to_0based_id[ids_to_keep] = torch.arange(n_ids_to_keep)
+        id_to_0based_id[ids_to_keep] = torch.arange(n_ids_to_keep, device=device)
         return id_to_0based_id[original_ids]
 
 

hyperbench/data/loader.py

@@ -97,13 +97,13 @@ def __batch_node_features(self, batch: List[HData]) -> Tuple[Tensor, int]:
         Returns:
             Tensor: Concatenated node features with shape (total_nodes, num_features).
         """
-        per_sample_node_features = [data.x for data in batch]
+        per_sample_x = [data.x for data in batch]
 
         # Stack all nodes along the node dimension from all samples into a single tensor
-        batched_node_features = torch.cat(per_sample_node_features, dim=0)
-        total_nodes = batched_node_features.size(0)
+        batched_x = torch.cat(per_sample_x, dim=0)
+        total_nodes = batched_x.size(0)
 
-        return batched_node_features, total_nodes
+        return batched_x, total_nodes
 
     def __batch_edges(self, batch: List[HData]) -> Tuple[Tensor, Optional[Tensor], int]:
         """Batches hyperedge indices and attributes, adjusting indices for concatenated nodes.

hyperbench/tests/data/dataset_test.py

@@ -727,7 +727,7 @@ class TestDataset(Dataset):
 
 
 def test_process_with_no_node_attributes_fallback():
-    """Test process() falls back to torch zeros when no numeric attributes."""
+    """Test process() falls back to torch ones when no node features."""
     mock_hypergraph = HIFHypergraph(
         network_type="undirected",
         nodes=[
@@ -746,7 +746,7 @@ class TestDataset(Dataset):
         dataset = TestDataset()
 
         assert dataset.hdata.x.shape == (2, 1)
-        assert torch.allclose(dataset.hdata.x, torch.tensor([[0.0], [0.0]]))
+        assert torch.allclose(dataset.hdata.x, torch.tensor([[1.0], [1.0]]))
 
 
 def test_process_with_single_node_attribute():