diff --git a/linear_operator/operators/kronecker_product_added_diag_linear_operator.py b/linear_operator/operators/kronecker_product_added_diag_linear_operator.py index e1685f74..60c43f88 100644 --- a/linear_operator/operators/kronecker_product_added_diag_linear_operator.py +++ b/linear_operator/operators/kronecker_product_added_diag_linear_operator.py @@ -84,7 +84,11 @@ 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() @@ -92,38 +96,17 @@ def _logdet( 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] diff --git a/test/operators/test_kronecker_product_added_diag_linear_operator.py b/test/operators/test_kronecker_product_added_diag_linear_operator.py index b055088b..0b820e34 100644 --- a/test/operators/test_kronecker_product_added_diag_linear_operator.py +++ b/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()