Update core/graph_mamba.py

core/graph_mamba.py

@@ -6,7 +6,7 @@ from .graph_sequencer import GraphSequencer, PositionalEncoder
 class GraphMamba(nn.Module):
     """
     Production Graph-Mamba model
-    Dynamically handles any graph size and structure
+    Device-safe implementation with dynamic handling
     """
     
     def __init__(self, config):
@@ -48,14 +48,22 @@ class GraphMamba(nn.Module):
         # Graph sequencer
         self.sequencer = GraphSequencer()
         
-    def _init_input_proj(self, input_dim):
+        # Classification head (for demo)
+        self.classifier = None
+        
+    def _init_input_proj(self, input_dim, device):
         """Initialize input projection dynamically"""
         if self.input_proj is None:
-            self.input_proj = nn.Linear(input_dim, self.d_model)
+            self.input_proj = nn.Linear(input_dim, self.d_model).to(device)
+            
+    def _init_classifier(self, num_classes, device):
+        """Initialize classifier dynamically"""
+        if self.classifier is None:
+            self.classifier = nn.Linear(self.d_model, num_classes).to(device)
             
     def forward(self, x, edge_index, batch=None):
         """
-        Forward pass with dynamic graph handling
+        Forward pass with device-safe handling
         
         Args:
             x: Node features (num_nodes, input_dim) 
@@ -64,9 +72,13 @@ class GraphMamba(nn.Module):
         """
         num_nodes = x.size(0)
         input_dim = x.size(1)
+        device = x.device
+        
+        # Move all components to correct device
+        self.to(device)
         
         # Initialize input projection if needed
-        self._init_input_proj(input_dim)
+        self._init_input_proj(input_dim, device)
         
         # Project input features
         h = self.input_proj(x)  # (num_nodes, d_model)
@@ -81,22 +93,31 @@ class GraphMamba(nn.Module):
         return h
     
     def _process_single_graph(self, h, edge_index):
-        """Process a single graph"""
+        """Process a single graph - device safe"""
         num_nodes = h.size(0)
+        device = h.device
+        
+        # Ensure edge_index is on correct device
+        edge_index = edge_index.to(device)
         
         # Get ordering
-        if self.ordering_strategy == "multi_view":
-            # Use BFS as primary for now (can be extended)
-            order = self.sequencer.bfs_ordering(edge_index, num_nodes)
-        elif self.ordering_strategy == "spectral":
+        if self.ordering_strategy == "spectral":
             order = self.sequencer.spectral_ordering(edge_index, num_nodes)
         elif self.ordering_strategy == "degree":
             order = self.sequencer.degree_ordering(edge_index, num_nodes)
+        elif self.ordering_strategy == "community":
+            order = self.sequencer.community_ordering(edge_index, num_nodes)
         else:  # default to BFS
             order = self.sequencer.bfs_ordering(edge_index, num_nodes)
         
+        # Ensure order is on correct device
+        order = order.to(device)
+        
         # Add positional encoding
         seq_pos, distances = self.pos_encoder.encode_positions(h, edge_index, order)
+        seq_pos = seq_pos.to(device)
+        distances = distances.to(device)
+        
         pos_features = torch.cat([seq_pos, distances], dim=1)  # (num_nodes, 11)
         pos_embed = self.pos_embed(pos_features)
         
@@ -119,7 +140,11 @@ class GraphMamba(nn.Module):
         return h_final
     
     def _process_batch(self, h, edge_index, batch):
-        """Process batched graphs"""
+        """Process batched graphs - device safe"""
+        device = h.device
+        batch = batch.to(device)
+        edge_index = edge_index.to(device)
+        
         batch_size = batch.max().item() + 1
         outputs = []
         
@@ -132,11 +157,15 @@ class GraphMamba(nn.Module):
             edge_mask = mask[edge_index[0]] & mask[edge_index[1]]
             batch_edges = edge_index[:, edge_mask]
             
-            # Reindex edges to local indices
-            node_indices = torch.where(mask)[0]
-            node_map = torch.zeros(h.size(0), dtype=torch.long, device=h.device)
-            node_map[node_indices] = torch.arange(batch_h.size(0), device=h.device)
-            batch_edges_local = node_map[batch_edges]
+            if batch_edges.shape[1] > 0:
+                # Reindex edges to local indices
+                node_indices = torch.where(mask)[0]
+                node_map = torch.zeros(h.size(0), dtype=torch.long, device=device)
+                node_map[node_indices] = torch.arange(batch_h.size(0), device=device)
+                batch_edges_local = node_map[batch_edges]
+            else:
+                # Empty graph
+                batch_edges_local = torch.empty((2, 0), dtype=torch.long, device=device)
             
             # Process subgraph
             batch_output = self._process_single_graph(batch_h, batch_edges_local)
@@ -144,7 +173,6 @@ class GraphMamba(nn.Module):
         
         # Reconstruct full batch
         h_out = torch.zeros_like(h)
-        start_idx = 0
         for b, output in enumerate(outputs):
             mask = batch == b
             h_out[mask] = output
@@ -157,6 +185,18 @@ class GraphMamba(nn.Module):
             # Single graph - mean pooling
             return h.mean(dim=0, keepdim=True)
         else:
-            # Batched graphs
-            from torch_geometric.nn import global_mean_pool
-            return global_mean_pool(h, batch)
+            # Batched graphs - manual pooling to avoid dependencies
+            device = h.device
+            batch = batch.to(device)
+            batch_size = batch.max().item() + 1
+            
+            graph_embeddings = []
+            for b in range(batch_size):
+                mask = batch == b
+                if mask.any():
+                    graph_emb = h[mask].mean(dim=0)
+                    graph_embeddings.append(graph_emb)
+                else:
+                    graph_embeddings.append(torch.zeros(h.size(1), device=device))
+            
+            return torch.stack(graph_embeddings)