Simplify batched Hurwitz operations (#54)

Author: lezcano

geotorch/hurwitz.py

@@ -1,5 +1,4 @@
 import torch
-import math
 from .psd import PSD
 from .skew import Skew
 from .product import ProductManifold
@@ -89,16 +88,8 @@ def submersion_inv(self, A: torch.Tensor, check_in_manifold=True, rho=1, tol=1e-
             A_shifted = A + self.alpha * self.In
             A_shifted_T = A_shifted.mT.contiguous()
 
-            # Batched Kronecker product using vmap for cleaner implementation
-            def _single_kron_sum(A_single):
-                I_single = torch.eye(self.n, device=A.device, dtype=A.dtype)
-                return torch.kron(I_single, A_single) + torch.kron(A_single, I_single)
-
-            flat_batch_size = math.prod(self.tensorial_size)
-            A_flat = A_shifted_T.reshape(flat_batch_size, self.n, self.n)
-
-            M_flat = torch.vmap(_single_kron_sum)(A_flat)
-            M = M_flat.reshape(*self.tensorial_size, self.n * self.n, self.n * self.n)
+            identity = torch.eye(self.n, device=A.device, dtype=A.dtype)
+            M = torch.kron(identity, A_shifted_T) + torch.kron(A_shifted_T, identity)
 
             Q = rho * self.In
             flat_Q = Q.flatten(-2, -1)
@@ -157,7 +148,7 @@ def sample(self, init_=torch.nn.init.xavier_normal_):
         with torch.no_grad():
             P = self[0].sample(init_) + self.In
             Q = self[1].sample(init_) + self.In
-            S = self[2](init_(torch.empty_like(P)))
+            S = self[2].sample(init_)
 
             A = self.submersion(Q, P, S)
             return A

test/test_hurwitz.py

@@ -109,6 +109,17 @@ def test_batch_submersion_inv(self):
 
         self.assertTrue(torch.allclose(A_batch, A_reconstructed, atol=1e-4))
 
+    def test_multi_batch_submersion_inv(self):
+        size = (2, 3, 2, 2)
+        hurwitz_batch = Hurwitz(size).double()
+        A_batch = hurwitz_batch.sample()
+        alpha = float(get_lyap_exp(A_batch)) / 2
+        inverse = Hurwitz(size, alpha=alpha).double()
+        Q, P_inv, S = inverse.submersion_inv(A_batch)
+        A_reconstructed = inverse.submersion(Q, P_inv, S)
+
+        self.assertTrue(torch.allclose(A_batch, A_reconstructed, atol=1e-4))
+
     def test_register_hurwitz(self):
         layer = torch.nn.Linear(self.size[-2], self.size[-1])
         geotorch.alpha_stable(layer, "weight", alpha=0.5)