diff --git a/singd/structures/base.py b/singd/structures/base.py index f8a9412..5408e69 100644 --- a/singd/structures/base.py +++ b/singd/structures/base.py @@ -251,23 +251,6 @@ def from_inner(self, X: Union[Tensor, None] = None) -> StructuredMatrix: S_dense = self.to_dense().T if X is None else self.rmatmat(X) return self.from_dense(supported_matmul(S_dense, S_dense.T)) - # NOTE This operation should be removed long-term as implementing IF-KFAC - # with `from_inner` is more efficient. For now, it will exist as it makes - # integrating this interface into existing implementations of sparse IF-KFAC - # easier, as they have access to the input/gradient covariance matrices. - def from_inner2(self, XXT: Tensor) -> StructuredMatrix: - """Extract the represented structure from ``self.T @ XXT @ self``. - - Args: - XXT: 2d square symmetric matrix. - - Returns: - The structured matrix extracted from ``self.T @ XXT @ self``. - """ - self._warn_naive_implementation("from_inner2") - dense = self.to_dense() - return self.from_dense(supported_matmul(dense.T, XXT, dense)) - def trace(self) -> Tensor: """Compute the trace of the represented matrix. diff --git a/singd/structures/diagonal.py b/singd/structures/diagonal.py index ac16b53..32a0ae7 100644 --- a/singd/structures/diagonal.py +++ b/singd/structures/diagonal.py @@ -133,18 +133,6 @@ def from_inner(self, X: Union[Tensor, None] = None) -> DiagonalMatrix: mat_diag *= (X**2).sum(1) return DiagonalMatrix(mat_diag) - def from_inner2(self, XXT: Tensor) -> StructuredMatrix: - """Represent the matrix diagonal of ``self.T @ XXT @ self``. - - Args: - XXT: 2d square symmetric matrix. - - Returns: - A ``DiagonalMatrix`` representing matrix diagonal of - ``self.T @ XXT @ self``. - """ - return DiagonalMatrix(self._mat_diag**2 * XXT.diag()) - def trace(self) -> Tensor: """Compute the trace of the represented matrix. diff --git a/singd/structures/triltoeplitz.py b/singd/structures/triltoeplitz.py index 650fd1b..368411a 100644 --- a/singd/structures/triltoeplitz.py +++ b/singd/structures/triltoeplitz.py @@ -26,7 +26,6 @@ class TrilToeplitzMatrix(StructuredMatrix): WARN_NAIVE_EXCEPTIONS = { # hard to leverage structure for efficient implementation "from_inner", - "from_inner2", } def __init__(self, diag_consts: Tensor) -> None: diff --git a/singd/structures/triutoeplitz.py b/singd/structures/triutoeplitz.py index a657fe8..b987249 100644 --- a/singd/structures/triutoeplitz.py +++ b/singd/structures/triutoeplitz.py @@ -26,7 +26,6 @@ class TriuToeplitzMatrix(StructuredMatrix): WARN_NAIVE_EXCEPTIONS = { # hard to leverage structure for efficient implementation "from_inner", - "from_inner2", } def __init__(self, diag_consts: Tensor) -> None: diff --git a/test/structures/utils.py b/test/structures/utils.py index 2c4ce4c..bf4747b 100644 --- a/test/structures/utils.py +++ b/test/structures/utils.py @@ -209,33 +209,6 @@ def _test_from_inner( ) -def _test_from_inner2( - sym_mat: Tensor, - structured_matrix_cls: Type[StructuredMatrix], - project: Callable[[Tensor], Tensor], - XXT: Tensor, -): - """Test `from_inner2` method of any child of `StructuredMatrix`. - - Args: - sym_mat: A symmetric dense matrix which will be converted into a structured - matrix. - structured_matrix_cls: The class of the structured matrix into which `sym_mat` - will be converted. - project: A function which converts an arbitrary symmetric dense matrix into a - dense matrix of the tested structure. Used to establish the ground truth. - XXT: An symmetric PSD matrix that will be passed to `from_inner2`. - """ - truth = project(supported_matmul(project(sym_mat).T, XXT, project(sym_mat))) - sym_mat_structured = structured_matrix_cls.from_dense(sym_mat) - report_nonclose( - truth, - sym_mat_structured.from_inner2(XXT).to_dense(), - rtol=1e-2 if is_half_precision(sym_mat.dtype) else 1e-5, - atol=1e-4 if is_half_precision(sym_mat.dtype) else 1e-7, - ) - - def _test_zeros( structured_matrix_cls: Type[StructuredMatrix], dim: int, @@ -454,23 +427,6 @@ def test_from_inner(self, dev: device, dtype: torch.dtype): X = rand((dim, 2 * dim), device=dev, dtype=dtype) _test_from_inner(sym_mat, self.STRUCTURED_MATRIX_CLS, self.project, X) - @mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) - @mark.parametrize("dev", DEVICES, ids=DEVICE_IDS) - def test_from_inner2(self, dev: device, dtype: torch.dtype): - """Test structure extraction after self-inner product w/ intermediate matrix. - - Args: - dev: The device on which to run the test. - dtype: The data type of the matrices. - """ - for dim in self.DIMS: - manual_seed(0) - - sym_mat = symmetrize(rand((dim, dim), device=dev, dtype=dtype)) - X = rand((dim, 2 * dim), device=dev, dtype=dtype) - XXT = supported_matmul(X, X.T) - _test_from_inner2(sym_mat, self.STRUCTURED_MATRIX_CLS, self.project, XXT) - @mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS) @mark.parametrize("dev", DEVICES, ids=DEVICE_IDS) def test_eye(self, dev: device, dtype: torch.dtype):