test fix

cornellius-gp · Nov 20, 2023 · 0d0e68f · 0d0e68f
1 parent d22b2be
commit 0d0e68f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/test/operators/test_diag_linear_operator.py b/test/operators/test_diag_linear_operator.py
@@ -11,6 +11,16 @@
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase
 
 
+def kron_diag(*lts):
+    """Compute diagonal of a KroneckerProductLinearOperator from the diagonals of the constituiting tensors"""
+    lead_diag = lts[0].diagonal(dim1=-1, dim2=-2)
+    if len(lts) == 1:  # base case:
+        return lead_diag
+    trail_diag = kron_diag(*lts[1:])
+    diag = lead_diag.unsqueeze(-2) * trail_diag.unsqueeze(-1)
+    return diag.mT.reshape(*diag.shape[:-2], -1)
+
+
 class TestDiagLinearOperator(LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True

diff --git a/test/operators/test_kronecker_product_linear_operator.py b/test/operators/test_kronecker_product_linear_operator.py
@@ -23,16 +23,6 @@ def kron(a, b):
     return torch.cat(res, -2)
 
 
-def kron_diag(*lts):
-    """Compute diagonal of a KroneckerProductLinearOperator from the diagonals of the constituiting tensors"""
-    lead_diag = lts[0].diagonal(dim1=-1, dim2=-2)
-    if len(lts) == 1:  # base case:
-        return lead_diag
-    trail_diag = kron_diag(*lts[1:])
-    diag = lead_diag.unsqueeze(-2) * trail_diag.unsqueeze(-1)
-    return diag.mT.reshape(*diag.shape[:-2], -1)
-
-
 class TestKroneckerProductLinearOperator(LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_call_lanczos = True