Fix head divisibility validation, benchmark data format, and target computation bugs

- Add embed_dim/feature_dim divisibility checks in attention modules
- Auto-select largest valid head count in IntradayGNN when num_features isn't divisible by 8
- Fix benchmark to pass PyG Data objects instead of tuples to model
- Fix target computation to use precomputed next trading date and filter valid tickers
- Add missing device argument to autocast calls
- Remove stale global app_state reference in dashboard API
- Clean up various formatting (line wrapping, whitespace)

src/amd/optimizations.py

@@ -165,6 +165,10 @@ class AMDSparseAttention(nn.Module):
 
     def __init__(self, embed_dim, num_heads, dropout=0.1):
         super().__init__()
+        if embed_dim % num_heads != 0:
+            raise ValueError(
+                f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})"
+            )
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         self.head_dim = embed_dim // num_heads
@@ -291,7 +295,9 @@ class AMDGATConv(MessagePassing):
         alpha = (alpha_src, alpha_dst)
 
         # Propagate — pass edge_attr so correlation weights reach message()
-        out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size)
+        out = self.propagate(
+            edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size
+        )
 
         # Concatenate or average heads
         if self.concat:

src/data/pipeline.py

@@ -309,7 +309,9 @@ class StockDataPipeline:
                     "high": row["High"],
                     "low": row["Low"],
                     "close": row["Close"],
-                    "adj_close": row["Close"],  # auto_adjust=True, 'Close' is already adjusted
+                    "adj_close": row[
+                        "Close"
+                    ],  # auto_adjust=True, 'Close' is already adjusted
                     "volume": row["Volume"],
                 }
             )
@@ -581,9 +583,7 @@ class StockDataPipeline:
         if len(seq_arr) > config.SEQUENCE_LENGTH:
             seq_arr = seq_arr[-config.SEQUENCE_LENGTH :]
         elif len(seq_arr) < config.SEQUENCE_LENGTH:
-            pad = np.zeros(
+            pad = np.zeros((config.SEQUENCE_LENGTH - len(seq_arr), 5), dtype=np.float32)
-                (config.SEQUENCE_LENGTH - len(seq_arr), 5), dtype=np.float32
-            )
             seq_arr = np.vstack([pad, seq_arr])
 
         # News features (current date, broadcast across all timesteps)
@@ -710,9 +710,7 @@ class StockDataPipeline:
                 continue
 
             # Pre-fetch PIT data once per ticker (avoids O(n²) repeated calls)
-            pit_cache = {
+            pit_cache = {t: self.get_point_in_time_data(t, date) for t in valid_tickers}
-                t: self.get_point_in_time_data(t, date) for t in valid_tickers
-            }
 
             # Build 3-D node feature matrix
             node_sequences: List[np.ndarray] = []
@@ -834,16 +832,33 @@ class StockDataPipeline:
             # Compute once per day — stock universe and PIT data don't change intraday
             pd_date = pd.Timestamp(date)
             current_tickers = [
-                t for t in tickers
+                t
+                for t in tickers
                 if t in self.price_data
                 and not self.price_data[t].empty
-                and self.price_data[t].index[0] <= pd_date <= self.price_data[t].index[-1]
+                and self.price_data[t].index[0]
+                <= pd_date
+                <= self.price_data[t].index[-1]
             ]
             if not current_tickers:
                 continue
 
             date_dt = datetime.strptime(date, "%Y-%m-%d")
-            pit_cache = {t: self.get_point_in_time_data(t, date_dt) for t in current_tickers}
+            pit_cache = {
+                t: self.get_point_in_time_data(t, date_dt) for t in current_tickers
+            }
+
+            # Determine next trading date once per day
+            next_date = self._next_trading_date(pd_date)
+            if next_date is None:
+                continue
+
+            # Pre-filter tickers that have target data on the next trading date
+            target_valid_tickers = [
+                t for t in current_tickers if next_date in self.price_data[t].index
+            ]
+            if not target_valid_tickers:
+                continue
 
             # Get all timestamps for this trading day
             timestamps = generate_intraday_timestamps(date)
@@ -856,7 +871,7 @@ class StockDataPipeline:
                 sequence_features = []
                 valid_tickers = []
 
-                for ticker in current_tickers:
+                for ticker in target_valid_tickers:
                     # Check memory before processing ticker
                     if not self.memory_manager.ensure_memory(50 * 1024**2):  # 50MB
                         logger.warning(f"Skipping {ticker} due to memory constraints")
@@ -881,8 +896,7 @@ class StockDataPipeline:
                         if isinstance(daily_data, pd.Series):
                             # If only one day of data, create a sequence with the same values
                             features = [
-                                daily_data["Close"] / daily_data["Open"]
+                                daily_data["Close"] / daily_data["Open"] - 1,  # Return
-                                - 1,  # Return
                                 0.2,  # Volatility (placeholder)
                                 0.0,  # Momentum (placeholder)
                                 np.log(daily_data["Volume"] + 1),  # Log volume
@@ -980,26 +994,14 @@ class StockDataPipeline:
                 )
 
                 # Create target (returns over prediction horizon)
-                y = []
+                y_vals = []
                 for ticker in valid_tickers:
-                    # For intraday, we would predict the next few minutes
+                    current_price = self.price_data[ticker].loc[date]["Close"]
-                    # For this example, we'll predict the next day's return
+                    future_price = self.price_data[ticker].loc[next_date]["Close"]
-                    next_date = datetime.strptime(date, "%Y-%m-%d") + timedelta(days=1)
+                    ret = future_price / current_price - 1
-                    if (
+                    y_vals.append(ret)
-                        ticker in self.price_data
-                        and next_date.strftime("%Y-%m-%d")
-                        in self.price_data[ticker].index
-                    ):
-                        current_price = self.price_data[ticker].loc[date]["Close"]
-                        future_price = self.price_data[ticker].loc[
-                            next_date.strftime("%Y-%m-%d")
-                        ]["Close"]
-                        ret = future_price / current_price - 1
-                        y.append(ret)
-                    else:
-                        y.append(0)  # Default value
 
-                y = torch.tensor(y, dtype=torch.float32).unsqueeze(1)
+                y = torch.tensor(y_vals, dtype=torch.float32).unsqueeze(1)
 
                 # Create Data object
                 data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y)

src/models/gnn_model.py

@@ -16,6 +16,10 @@ class TemporalAttention(nn.Module):
 
     def __init__(self, feature_dim: int, num_heads: int = 8, dropout: float = 0.1):
         super().__init__()
+        if feature_dim % num_heads != 0:
+            raise ValueError(
+                f"feature_dim ({feature_dim}) must be divisible by num_heads ({num_heads})"
+            )
         self.feature_dim = feature_dim
         self.num_heads = num_heads
         self.head_dim = feature_dim // num_heads
@@ -308,7 +312,9 @@ class CorporateActionAwareGNN(nn.Module):
 
         price_attended = self.temporal_attention(self.price_processor(price_features))
         news_attended = self.temporal_attention(self.news_processor(news_features))
-        social_attended = self.temporal_attention(self.social_processor(social_features))
+        social_attended = self.temporal_attention(
+            self.social_processor(social_features)
+        )
 
         alternative_features = torch.cat(
             [price_attended, news_attended, social_attended], dim=1

src/models/intraday_gnn.py

@@ -24,7 +24,12 @@ class IntradayGNN(nn.Module):
         self.num_features = num_features
 
         # Temporal attention for sequence processing
-        self.temporal_attention = TemporalAttention(num_features, num_heads=8)
+        # num_features must be divisible by num_heads; 14 is not divisible by 8.
+        # Pick the largest divisor of num_features up to 8.
+        num_heads = max(
+            h for h in range(1, min(9, num_features + 1)) if num_features % h == 0
+        )
+        self.temporal_attention = TemporalAttention(num_features, num_heads=num_heads)
 
         # Feature processing modules with AMD optimizations
         self.feature_processor = nn.Sequential(
@@ -161,8 +166,8 @@ class IntradayGNN(nn.Module):
         lstm_out = lstm_out.squeeze(1)  # (N, HIDDEN_CHANNELS)
 
         # Per-feature gating: sigmoid weights scale each channel independently
-        attention_weights = self.attention(lstm_out)   # (N, HIDDEN_CHANNELS)
+        attention_weights = self.attention(lstm_out)  # (N, HIDDEN_CHANNELS)
-        attended = lstm_out * attention_weights         # (N, HIDDEN_CHANNELS)
+        attended = lstm_out * attention_weights  # (N, HIDDEN_CHANNELS)
 
         return self.linear(attended)
 

src/models/trainer.py

@@ -143,7 +143,9 @@ class GNNTrainer:
                     self.memory_manager.empty_cache()
                     continue
 
-            epoch_train_loss = epoch_train_loss / num_train_batches if num_train_batches else 0.0
+            epoch_train_loss = (
+                epoch_train_loss / num_train_batches if num_train_batches else 0.0
+            )
             self.train_losses.append(epoch_train_loss)
 
             # Validation
@@ -185,7 +187,9 @@ class GNNTrainer:
                         self.memory_manager.empty_cache()
                         continue
 
-            epoch_val_loss = epoch_val_loss / num_val_batches if num_val_batches else 0.0
+            epoch_val_loss = (
+                epoch_val_loss / num_val_batches if num_val_batches else 0.0
+            )
             self.val_losses.append(epoch_val_loss)
 
             # Update learning rate scheduler
@@ -234,6 +238,7 @@ class GNNTrainer:
             data = data.to(self.device, non_blocking=config.PIN_MEMORY)
 
             with autocast(
+                config.DEVICE,
                 enabled=config.MIXED_PRECISION,
                 dtype=self.amd_optimizer.get_precision_dtype(),
             ):
@@ -343,6 +348,7 @@ class GNNTrainer:
                 data = data.to(self.device, non_blocking=config.PIN_MEMORY)
 
                 with autocast(
+                    config.DEVICE,
                     enabled=config.MIXED_PRECISION,
                     dtype=self.amd_optimizer.get_precision_dtype(),
                 ):

src/web/api/dashboard.py

@@ -13,9 +13,6 @@ from src.web.services.state import AppState
 logger = logging.getLogger(__name__)
 router = APIRouter()
 
-# Reference to global app state (injected via module import in app.py)
-app_state: AppState = None  # type: ignore
-
 
 def _get_state() -> AppState:
     from src.web.app import app_state as _state

src/web/api/models_endpoints.py

@@ -8,6 +8,7 @@ from datetime import datetime
 from typing import Dict
 
 from fastapi import APIRouter
+from torch_geometric.data import Data
 
 from config import config
 
@@ -151,16 +152,17 @@ async def run_benchmark() -> Dict:
     edge_attr = torch.randn(num_edges, 1).to(config.DEVICE)
 
     # Warm-up
+    dummy_data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
     for _ in range(10):
         with torch.no_grad():
-            _ = model((x, edge_index, edge_attr))
+            _ = model(dummy_data)
 
     # Benchmark inference
     start = time.time()
     num_runs = 100
     for _ in range(num_runs):
         with torch.no_grad():
-            _ = model((x, edge_index, edge_attr))
+            _ = model(dummy_data)
     inference_time = (time.time() - start) / num_runs
 
     # Benchmark training
@@ -172,7 +174,7 @@ async def run_benchmark() -> Dict:
     start = time.time()
     for _ in range(num_runs):
         optimizer.zero_grad()
-        out = model((x, edge_index, edge_attr))
+        out = model(dummy_data)
         loss = criterion(out, y)
         loss.backward()
         optimizer.step()