Update test_cases.py

test_cases.py

@@ -265,15 +265,8 @@ class MagnitudeFlow(BaseFlow):
         # Gram matrix
         G = torch.bmm(a_geom.transpose(-2, -1), a_geom)  # [B, geom_dim, 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]
+        # Eigendecomposition — differentiable through torch.linalg.eigh
+        eigenvalues, _ = LA.eigh(G, method='torch')  # [B, geom_dim]
 
         # Magnitude spectrum: how energy distributes across modes
         magnitudes = eigenvalues.abs().sqrt()  # [B, geom_dim] — the ω spectrum
@@ -321,14 +314,8 @@ 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]
 
-        # 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]
+        # Eigendecomposition — the ω spectrum (differentiable via torch.linalg.eigh)
+        eigenvalues, eigenvectors = LA.eigh(G, method='torch')  # [B, gd], [B, gd, gd]
 
         # ω = √|λ|
         omega = eigenvalues.abs().sqrt()  # [B, gd]
@@ -360,40 +347,41 @@ class OrbitalFlow(BaseFlow):
 # ═══════════════════════════════════════════════════════════════════
 
 class AlignmentFlow(BaseFlow):
-    """SVD alignment flow via soft Procrustes rotation.
+    """SVD alignment flow via soft Procrustes rotation in projected space.
 
-    Computes attention-weighted anchor targets per query, then finds the
-    optimal rotation aligning queries toward those targets via SVD of
-    the cross-covariance matrix.
+    Projects to geom_dim, computes optimal rotation via SVD of the
+    cross-covariance in the small space, applies rotation, projects back.
     """
     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.geom_dim = min(n_anchors, 12)  # FL eigh sweet spot
+        self.anchor_proj = nn.Linear(d_model, self.geom_dim)
+        self.query_proj = nn.Linear(d_model, self.geom_dim)
+        self.geom_to_query = nn.Linear(self.geom_dim, d_model)
         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]
+        # Project to small geometric space
+        a_proj = self.anchor_proj(anchors)  # [B, k, geom_dim]
+        q_proj = self.query_proj(queries)   # [B, n, geom_dim]
 
-        # Attention-weighted anchors → per-query targets [B, n, d]
-        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1)) / math.sqrt(d)
+        # Attention-weighted anchors → per-query targets [B, n, geom_dim]
+        sim = torch.bmm(q_proj, a_proj.transpose(-2, -1)) / math.sqrt(self.geom_dim)
         weights = F.softmax(sim, dim=-1)  # [B, n, k]
-        targets = torch.bmm(weights, a_proj)  # [B, n, d]
+        targets = torch.bmm(weights, a_proj)  # [B, n, geom_dim]
 
-        # Cross-covariance: C = Q^T @ T, both [B, n, d] → C is [B, d, d]
+        # Cross-covariance in small space: [B, geom_dim, geom_dim]
         C = torch.bmm(q_proj.transpose(-2, -1), targets)
 
-        # SVD → optimal rotation (Procrustes)
-        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]
+        # SVD → optimal rotation via gram_eigh (differentiable, no in-place ops)
+        U, _, Vh = LA.svd(C, method='gram_eigh')
+        R = torch.bmm(U, Vh)  # [B, geom_dim, geom_dim]
 
-        # Soft rotation
-        q_rotated = torch.bmm(queries, R)
-        return queries + self.strength * (q_rotated - queries)
+        # Rotate queries in small space, project back to d_model
+        q_rotated = torch.bmm(q_proj, R)  # [B, n, geom_dim]
+        delta = self.geom_to_query(q_rotated - q_proj)  # [B, n, d]
+        return queries + self.strength * delta
 
 
 # ═════════════════════════════════════════��═════════════════════════