From e851bf24f423fb60d97121853729f73ed43755ea Mon Sep 17 00:00:00 2001 From: lezcano Date: Sun, 7 Jun 2026 11:14:21 +0100 Subject: [PATCH] Simplify batched Hurwitz operations --- geotorch/hurwitz.py | 15 +++------------ test/test_hurwitz.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/geotorch/hurwitz.py b/geotorch/hurwitz.py index ebd5f3a4..478a698f 100644 --- a/geotorch/hurwitz.py +++ b/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 diff --git a/test/test_hurwitz.py b/test/test_hurwitz.py index 078571eb..6776f136 100644 --- a/test/test_hurwitz.py +++ b/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)