Update test_cases.py

test_cases.py

@@ -33,6 +33,13 @@ import torch.nn.functional as F
 from torch import Tensor
 from typing import List, Optional, Tuple
 
+# Use geolip_core.linalg when available (FL eigh, Triton SVD, etc.)
+# Falls back to torch.linalg transparently
+try:
+    import geolip_core.linalg as LA
+except ImportError:
+    import torch.linalg as LA
+
 
 # ═══════════════════════════════════════════════════════════════════
 # Base Flow
@@ -255,10 +262,18 @@ class MagnitudeFlow(BaseFlow):
         # Project anchors to geometric space
         a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]
 
-        # Gram matrix eigenvalues → spectral profile
+        # Gram matrix
         G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, geom_dim, geom_dim]
-        # Use torch.linalg.eigh for now; swap to FL eigh in geolip.linalg
-        eigenvalues, _ = torch.linalg.eigh(G)  # [B, geom_dim]
+
+        # Eigenvectors under no_grad (FL eigh has in-place deflation)
+        # Eigenvalues recomputed via Rayleigh quotient for differentiability
+        with torch.no_grad():
+            _, V = LA.eigh(G)  # [B, gd, gd]
+        V = V.detach()
+
+        # Rayleigh quotient: λᵢ = vᵢᵀ G vᵢ — differentiable through G
+        GV = torch.bmm(G, V)  # [B, gd, gd]
+        eigenvalues = (V * GV).sum(dim=-2)  # [B, gd]
 
         # Magnitude spectrum: how energy distributes across modes
         magnitudes = eigenvalues.abs().sqrt()  # [B, geom_dim] — the ω spectrum
@@ -306,8 +321,14 @@ class OrbitalFlow(BaseFlow):
         a_geom = self.anchor_proj(anchors)  # [B, k, geom_dim]
         G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, gd, gd]
 
-        # Eigendecomposition — the ω spectrum
-        eigenvalues, eigenvectors = torch.linalg.eigh(G)  # [B, gd], [B, gd, gd]
+        # Eigenvectors under no_grad (FL eigh has in-place deflation)
+        with torch.no_grad():
+            _, eigenvectors = LA.eigh(G)  # [B, gd, gd]
+        eigenvectors = eigenvectors.detach()
+
+        # Rayleigh quotient: differentiable eigenvalues through G
+        GV = torch.bmm(G, eigenvectors)  # [B, gd, gd]
+        eigenvalues = (eigenvectors * GV).sum(dim=-2)  # [B, gd]
 
         # ω = √|λ|
         omega = eigenvalues.abs().sqrt()  # [B, gd]
@@ -341,29 +362,36 @@ class OrbitalFlow(BaseFlow):
 class AlignmentFlow(BaseFlow):
     """SVD alignment flow via soft Procrustes rotation.
 
-    Computes the optimal rotation aligning queries toward the anchor
-    geometry using SVD of the cross-covariance matrix. The rotation
-    is applied as a soft geometric bias.
+    Computes attention-weighted anchor targets per query, then finds the
+    optimal rotation aligning queries toward those targets via SVD of
+    the cross-covariance matrix.
     """
     def __init__(self, d_model: int, n_anchors: int):
         super().__init__(d_model, n_anchors, name='alignment')
         self.anchor_proj = nn.Linear(d_model, d_model)
         self.query_proj = nn.Linear(d_model, d_model)
-        self.strength = nn.Parameter(torch.tensor(0.1))  # learnable blend
+        self.strength = nn.Parameter(torch.tensor(0.1))
 
     def _flow(self, anchors, queries):
         B, n, d = queries.shape
         a_proj = self.anchor_proj(anchors)  # [B, k, d]
         q_proj = self.query_proj(queries)   # [B, n, d]
 
-        # Cross-covariance: C = Q^T A
-        C = torch.bmm(q_proj.transpose(-2, -1), a_proj)  # [B, d, d]
+        # Attention-weighted anchors → per-query targets [B, n, d]
+        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1)) / math.sqrt(d)
+        weights = F.softmax(sim, dim=-1)  # [B, n, k]
+        targets = torch.bmm(weights, a_proj)  # [B, n, d]
+
+        # Cross-covariance: C = Q^T @ T, both [B, n, d] → C is [B, d, d]
+        C = torch.bmm(q_proj.transpose(-2, -1), targets)
 
         # SVD → optimal rotation (Procrustes)
-        U, _, Vh = torch.linalg.svd(C)
-        R = torch.bmm(U, Vh)  # [B, d, d] rotation matrix
+        U, _, Vh = torch.linalg.svd(C)  # full d×d SVD — not through geolip.linalg.svd
+        # Note: geolip.linalg.svd is thin SVD for M≥N rectangular matrices.
+        # Cross-covariance C is square [B, d, d], use torch directly.
+        R = torch.bmm(U, Vh)  # [B, d, d]
 
-        # Apply soft rotation
+        # Soft rotation
         q_rotated = torch.bmm(queries, R)
         return queries + self.strength * (q_rotated - queries)