refactor: Move autogram._gramian_utils to _linalg.generalized_gramian (#627)

Author: PierreQuinton

src/torchjd/_linalg/__init__.py

@@ -1,3 +1,4 @@
+from ._generalized_gramian import flatten, movedim, reshape
 from ._gramian import compute_gramian, normalize, regularize
 from ._matrix import Matrix, PSDMatrix, PSDTensor, is_matrix, is_psd_matrix, is_psd_tensor
 
@@ -11,4 +12,7 @@
     "is_matrix",
     "is_psd_matrix",
     "is_psd_tensor",
+    "flatten",
+    "reshape",
+    "movedim",
 ]

src/torchjd/autogram/_gramian_utils.py …/torchjd/_linalg/_generalized_gramian.pysrc/torchjd/autogram/_gramian_utils.py renamed to src/torchjd/_linalg/_generalized_gramian.py src/torchjd/autogram/_gramian_utils.py renamed to src/torchjd/_linalg/_generalized_gramian.py

@@ -3,7 +3,7 @@
 
 from torch import Tensor
 
-from torchjd._linalg import PSDMatrix, PSDTensor
+from torchjd._linalg._matrix import PSDMatrix, PSDTensor
 
 
 def flatten(gramian: PSDTensor) -> PSDMatrix:

src/torchjd/aggregation/_flattening.py

@@ -1,8 +1,7 @@
 from torch import Tensor
 
-from torchjd._linalg import PSDTensor
+from torchjd._linalg import PSDTensor, flatten
 from torchjd.aggregation._weighting_bases import GeneralizedWeighting, Weighting
-from torchjd.autogram._gramian_utils import flatten
 
 
 class Flattening(GeneralizedWeighting):

src/torchjd/autogram/_engine.py

@@ -4,12 +4,11 @@
 from torch import Tensor, nn, vmap
 from torch.autograd.graph import get_gradient_edge
 
-from torchjd._linalg import PSDMatrix
+from torchjd._linalg import PSDMatrix, movedim, reshape
 
 from ._edge_registry import EdgeRegistry
 from ._gramian_accumulator import GramianAccumulator
 from ._gramian_computer import GramianComputer, JacobianBasedGramianComputerWithCrossTerms
-from ._gramian_utils import movedim, reshape
 from ._jacobian_computer import (
     AutogradJacobianComputer,
     FunctionalJacobianComputer,

tests/unit/autogram/test_engine.py

@@ -79,10 +79,9 @@
 )
 from utils.tensors import make_inputs_and_targets, ones_, randn_, zeros_
 
-from torchjd._linalg import PSDMatrix, compute_gramian
+from torchjd._linalg import PSDMatrix, compute_gramian, movedim, reshape
 from torchjd.aggregation import UPGradWeighting
 from torchjd.autogram._engine import Engine
-from torchjd.autogram._gramian_utils import movedim, reshape
 
 PARAMETRIZATIONS = [
     (ModuleFactory(OverlyNested), 32),

tests/unit/autogram/test_gramian_utils.py …/unit/linalg/test_generalized_gramian.pytests/unit/autogram/test_gramian_utils.py renamed to tests/unit/linalg/test_generalized_gramian.py tests/unit/autogram/test_gramian_utils.py renamed to tests/unit/linalg/test_generalized_gramian.py

@@ -3,8 +3,7 @@
 from utils.asserts import assert_is_psd_matrix, assert_is_psd_tensor
 from utils.tensors import randn_
 
-from torchjd._linalg import compute_gramian, is_psd_matrix
-from torchjd.autogram._gramian_utils import flatten, movedim, reshape
+from torchjd._linalg import compute_gramian, flatten, is_psd_matrix, movedim, reshape
 
 
 @mark.parametrize(

tests/utils/asserts.py

@@ -4,8 +4,7 @@
 from torch import Tensor
 from torch.testing import assert_close
 
-from torchjd._linalg import is_psd_matrix, is_psd_tensor
-from torchjd.autogram._gramian_utils import flatten
+from torchjd._linalg import flatten, is_psd_matrix, is_psd_tensor
 from torchjd.autojac._accumulation import is_tensor_with_jac