Fix log determinants with constant Kronecker diagonal noise

linear_operator/operators/kronecker_product_added_diag_linear_operator.py

@@ -84,46 +84,29 @@ def inv_quad_logdet(
     def _logdet(
         self: LinearOperator,  # shape: (*batch, N, N)
     ) -> Tensor:  # shape: (*batch)
-        if self._diag_is_constant:
+        if self._diag_is_constant or (
+            self.shape[-1] >= settings.max_cholesky_size.value()
+            and isinstance(self.diag_tensor, KroneckerProductDiagLinearOperator)
+            and all(isinstance(dt, ConstantDiagLinearOperator) for dt in self.diag_tensor.linear_ops)
+        ):
             # symeig requires computing the eigenvectors for it to be differentiable
             evals, _ = self.linear_op._symeig(eigenvectors=True)
             evals_plus_diag = evals + self.diag_tensor._diagonal()
             return torch.log(evals_plus_diag).sum(dim=-1)
         if self.shape[-1] >= settings.max_cholesky_size.value() and isinstance(
             self.diag_tensor, KroneckerProductDiagLinearOperator
         ):
-            # If the diagonal has the same Kronecker structure as the full matrix, with each factor being
-            # constant, wee can compute the logdet efficiently
-            if len(self.linear_op.linear_ops) == len(self.diag_tensor.linear_ops) and all(
-                isinstance(dt, ConstantDiagLinearOperator) for dt in self.diag_tensor.linear_ops
-            ):
-                # here the log determinant identity is |D + K| = | D| |I + D^{-1} K|
-                # as D is assumed to have constant components, we can look solely at the diag_values
-                diag_term = self.diag_tensor._diagonal().clamp(min=1e-7).log().sum(dim=-1)
-                # symeig requires computing the eigenvectors for it to be differentiable
-                evals, _ = self.linear_op._symeig(eigenvectors=True)
-                const_times_evals = KroneckerProductLinearOperator(
-                    *[ee * d.diag_values for ee, d in zip(evals.linear_ops, self.diag_tensor.linear_ops)]
+            # Use |K + D| = |D| |I + D^{-1}K|, symmetrizing the second
+            # matrix because its eigendecomposition must remain differentiable.
+            lt = self.linear_op
+            dlt = self.diag_tensor
+            if isinstance(lt, KroneckerProductAddedDiagLinearOperator):
+                raise NotImplementedError(
+                    "Log determinant for KroneckerProductAddedDiagLinearOperator + "
+                    "DiagLinearOperator not implemented."
                 )
-                first_term = (const_times_evals._diagonal() + 1).log().sum(dim=-1)
-                return diag_term + first_term
-
-            else:
-                # we use the same matrix determinant identity: |K + D| = |D| |I + D^{-1}K|
-                # but have to symmetrize the second matrix because torch.linalg.eig may not be
-                # completely differentiable.
-                lt = self.linear_op
-                dlt = self.diag_tensor
-                if isinstance(lt, KroneckerProductAddedDiagLinearOperator):
-                    raise NotImplementedError(
-                        "Log determinant for KroneckerProductAddedDiagLinearOperator + "
-                        "DiagLinearOperator not implemented."
-                    )
-                else:
-                    _, evals_plus_i, _ = _symmetrize_kpadlt_constructor(lt, dlt)
-
-                diag_term = self.diag_tensor.logdet()
-                return diag_term + evals_plus_i.logdet()
+            _, evals_plus_i, _ = _symmetrize_kpadlt_constructor(lt, dlt)
+            return self.diag_tensor.logdet() + evals_plus_i.logdet()
 
         return super().inv_quad_logdet(logdet=True)[1]
 

test/operators/test_kronecker_product_added_diag_linear_operator.py

@@ -158,5 +158,36 @@ def test_root_inv_decomposition_no_cholesky(self):
                 chol_mock.assert_not_called()
 
 
+class TestConstantKroneckerDiagonalLogdet(unittest.TestCase):
+    def test_logdet_and_noise_derivatives(self):
+        for dtype in (torch.float32, torch.float64):
+            for constants in ((2.0, 3.0), (0.0, 3.0), (-2.0, -3.0)):
+                with self.subTest(dtype=dtype, constants=constants):
+                    a = torch.tensor([[2.0, 0.2], [0.2, 3.0]], dtype=dtype)
+                    b = torch.tensor([[1.5, 0.3], [0.3, 2.5]], dtype=dtype)
+                    noise = torch.tensor(constants, dtype=dtype, requires_grad=True)
+                    kron = KroneckerProductLinearOperator(a, b)
+                    diag = KroneckerProductDiagLinearOperator(
+                        ConstantDiagLinearOperator(noise[:1], diag_shape=2),
+                        ConstantDiagLinearOperator(noise[1:], diag_shape=2),
+                    )
+                    linear_op = KroneckerProductAddedDiagLinearOperator(kron, diag)
+                    dense = torch.kron(a, b) + noise.prod() * torch.eye(4, dtype=dtype)
+                    with settings.max_cholesky_size(0):
+                        actual = linear_op.logdet()
+                    expected = torch.linalg.slogdet(dense).logabsdet
+                    torch.testing.assert_close(actual, expected)
+                    actual_grad = torch.autograd.grad(actual, noise, create_graph=True)[0]
+                    expected_grad = torch.autograd.grad(expected, noise, create_graph=True)[0]
+                    torch.testing.assert_close(actual_grad, expected_grad)
+                    actual_hessian = torch.stack(
+                        [torch.autograd.grad(g, noise, retain_graph=True)[0] for g in actual_grad]
+                    )
+                    expected_hessian = torch.stack(
+                        [torch.autograd.grad(g, noise, retain_graph=True)[0] for g in expected_grad]
+                    )
+                    torch.testing.assert_close(actual_hessian, expected_hessian)
+
+
 if __name__ == "__main__":
     unittest.main()