Can parametrize number of dimensions to contract in compute_gramian

Author: PierreQuinton

src/torchjd/_linalg/_gramian.py

@@ -1,18 +1,40 @@
-from typing import cast
+from typing import Literal, cast, overload
 
 import torch
 
-from ._matrix import GeneralizedMatrix, PSDMatrix
+from ._matrix import GeneralizedMatrix, PSDGeneralizedMatrix, PSDMatrix
 
 
+@overload
 def compute_gramian(matrix: GeneralizedMatrix) -> PSDMatrix:
+    pass
+
+
+@overload
+def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: Literal[-1]) -> PSDMatrix:
+    pass
+
+
+@overload
+def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: int) -> PSDGeneralizedMatrix:
+    pass
+
+
+def compute_gramian(matrix: GeneralizedMatrix, contracted_dims: int = -1) -> PSDGeneralizedMatrix:
     """
-    Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of a given matrix.
+    Computes the `Gramian matrix <https://en.wikipedia.org/wiki/Gram_matrix>`_ of the input.
+
+    `contracted_dims` specifies the number of trailing dimensions to contract. If negative,
+    it indicates the number of leading dimensions to preserve (e.g., ``-1`` preserves the
+    first dimension).
     """
 
-    indices = list(range(1, matrix.ndim))
-    gramian = torch.tensordot(matrix, matrix, dims=(indices, indices))
-    return cast(PSDMatrix, gramian)
+    contracted_dims = contracted_dims if 0 <= contracted_dims else contracted_dims + matrix.ndim
+    indices_source = list(range(matrix.ndim - contracted_dims))
+    indices_dest = list(range(matrix.ndim - 1, contracted_dims - 1, -1))
+    transposed_matrix = matrix.movedim(indices_source, indices_dest)
+    gramian = torch.tensordot(matrix, transposed_matrix, dims=contracted_dims)
+    return cast(PSDGeneralizedMatrix, gramian)
 
 
 def normalize(gramian: PSDMatrix, eps: float) -> PSDMatrix:

tests/unit/autogram/test_engine.py

@@ -68,7 +68,6 @@
     CloneParams,
     autograd_forward_backward,
     autogram_forward_backward,
-    compute_gramian,
     compute_gramian_with_autograd,
     forward_pass,
     make_mse_loss_fn,
@@ -79,6 +78,7 @@
 )
 from utils.tensors import make_inputs_and_targets, ones_, randn_, zeros_
 
+from torchjd._linalg import compute_gramian
 from torchjd.aggregation import UPGradWeighting
 from torchjd.autogram._engine import Engine
 from torchjd.autogram._gramian_utils import movedim, reshape
@@ -418,9 +418,9 @@ def test_compute_gramian_manual():
     weight_jacobian = zeros_([out_dims, model.weight.numel()])
     for j in range(out_dims):
         weight_jacobian[j, j * in_dims : (j + 1) * in_dims] = input
-    weight_gramian = compute_gramian(weight_jacobian)
+    weight_gramian = compute_gramian(weight_jacobian, 1)
     bias_jacobian = torch.diag(ones_(out_dims))
-    bias_gramian = compute_gramian(bias_jacobian)
+    bias_gramian = compute_gramian(bias_jacobian, 1)
     expected_gramian = weight_gramian + bias_gramian
 
     assert_close(gramian, expected_gramian)

tests/unit/autogram/test_gramian_utils.py

@@ -1,10 +1,9 @@
 from pytest import mark
 from torch.testing import assert_close
 from utils.asserts import assert_is_psd_generalized_matrix, assert_is_psd_matrix
-from utils.forward_backwards import compute_gramian
 from utils.tensors import randn_
 
-from torchjd._linalg import is_psd_matrix
+from torchjd._linalg import compute_gramian, is_psd_matrix
 from torchjd.autogram._gramian_utils import flatten, movedim, reshape
 
 
@@ -33,8 +32,8 @@ def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]
     original_matrix = randn_(original_shape + [2])
     target_matrix = original_matrix.reshape(target_shape + [2])
 
-    original_gramian = compute_gramian(original_matrix)
-    target_gramian = compute_gramian(target_matrix)
+    original_gramian = compute_gramian(original_matrix, 1)
+    target_gramian = compute_gramian(target_matrix, 1)
 
     reshaped_gramian = reshape(original_gramian, target_shape)
 
@@ -58,7 +57,7 @@ def test_reshape_equivarience(original_shape: list[int], target_shape: list[int]
 )
 def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
     matrix = randn_(original_shape + [2])
-    gramian = compute_gramian(matrix)
+    gramian = compute_gramian(matrix, 1)
     reshaped_gramian = reshape(gramian, target_shape)
     assert_is_psd_generalized_matrix(reshaped_gramian, atol=1e-04, rtol=0.0)
 
@@ -75,7 +74,7 @@ def test_reshape_yields_psd(original_shape: list[int], target_shape: list[int]):
 )
 def test_flatten_yields_matrix(shape: list[int]):
     matrix = randn_(shape + [2])
-    gramian = compute_gramian(matrix)
+    gramian = compute_gramian(matrix, 1)
     flattened_gramian = flatten(gramian)
     assert is_psd_matrix(flattened_gramian)
 
@@ -92,7 +91,7 @@ def test_flatten_yields_matrix(shape: list[int]):
 )
 def test_flatten_yields_psd(shape: list[int]):
     matrix = randn_(shape + [2])
-    gramian = compute_gramian(matrix)
+    gramian = compute_gramian(matrix, 1)
     flattened_gramian = flatten(gramian)
     assert_is_psd_matrix(flattened_gramian, atol=1e-04, rtol=0.0)
 
@@ -121,8 +120,8 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
     original_matrix = randn_(shape + [2])
     target_matrix = original_matrix.movedim(source, destination)
 
-    original_gramian = compute_gramian(original_matrix)
-    target_gramian = compute_gramian(target_matrix)
+    original_gramian = compute_gramian(original_matrix, 1)
+    target_gramian = compute_gramian(target_matrix, 1)
 
     moveddim_gramian = movedim(original_gramian, source, destination)
 
@@ -149,6 +148,6 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
 )
 def test_movedim_yields_psd(shape: list[int], source: list[int], destination: list[int]):
     matrix = randn_(shape + [2])
-    gramian = compute_gramian(matrix)
+    gramian = compute_gramian(matrix, 1)
     moveddim_gramian = movedim(gramian, source, destination)
     assert_is_psd_generalized_matrix(moveddim_gramian)

tests/unit/linalg/test_gramian.py

@@ -2,8 +2,7 @@
 from utils.asserts import assert_is_psd_matrix
 from utils.tensors import randn_
 
-from torchjd._linalg import compute_gramian, is_generalized_matrix, is_matrix
-from torchjd._linalg._gramian import normalize, regularize
+from torchjd._linalg import compute_gramian, is_generalized_matrix, is_matrix, normalize, regularize
 
 
 @mark.parametrize(

tests/utils/forward_backwards.py

@@ -142,14 +142,6 @@ def get_vjp(grad_outputs: Tensor) -> list[Tensor]:
     return gramian
 
 
-def compute_gramian(matrix: Tensor) -> PSDGeneralizedMatrix:
-    """Contracts the last dimension of matrix to make it into a Gramian."""
-
-    indices = list(range(matrix.ndim))
-    transposed_matrix = matrix.movedim(indices, indices[::-1])
-    return torch.tensordot(matrix, transposed_matrix, dims=([-1], [0]))
-
-
 class CloneParams:
     """
     ContextManager enabling the computation of per-usage gradients.