Rename DiagonalSparseTensor to StructuredSparseTensor

Author: ValerianRey

src/torchjd/autogram/_engine.py

@@ -4,7 +4,7 @@
 from torch import Tensor, nn, vmap
 from torch.autograd.graph import get_gradient_edge
 
-from torchjd.sparse import make_dst
+from torchjd.sparse import make_sst
 
 from ._edge_registry import EdgeRegistry
 from ._gramian_accumulator import GramianAccumulator
@@ -176,7 +176,7 @@ def differentiation(_grad_output: Tensor) -> tuple[Tensor, ...]:
 
             output_dims = list(range(output.ndim))
             v_to_ps = [[dim] for dim in output_dims * 2]
-            jac_output = make_dst(torch.ones_like(output), v_to_ps)
+            jac_output = make_sst(torch.ones_like(output), v_to_ps)
 
             vmapped_diff = differentiation
             for _ in output_dims:

src/torchjd/sparse/__init__.py

@@ -1,3 +1,3 @@
 # Need to import this to execute the code inside and thus to override the functions
 from . import _aten_function_overrides
-from ._diagonal_sparse_tensor import DiagonalSparseTensor, make_dst
+from ._structured_sparse_tensor import StructuredSparseTensor, make_sst

src/torchjd/sparse/_aten_function_overrides/backward.py

@@ -1,36 +1,36 @@
 from torch import Tensor
 from torch.ops import aten  # type: ignore
 
-from torchjd.sparse import DiagonalSparseTensor
+from torchjd.sparse._structured_sparse_tensor import StructuredSparseTensor
 
 
-@DiagonalSparseTensor.implements(aten.threshold_backward.default)
+@StructuredSparseTensor.implements(aten.threshold_backward.default)
 def threshold_backward_default(
-    grad_output: DiagonalSparseTensor, self: Tensor, threshold
-) -> DiagonalSparseTensor:
+    grad_output: StructuredSparseTensor, self: Tensor, threshold
+) -> StructuredSparseTensor:
     new_physical = aten.threshold_backward.default(grad_output.physical, self, threshold)
 
-    return DiagonalSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)
 
 
-@DiagonalSparseTensor.implements(aten.hardtanh_backward.default)
+@StructuredSparseTensor.implements(aten.hardtanh_backward.default)
 def hardtanh_backward_default(
-    grad_output: DiagonalSparseTensor,
+    grad_output: StructuredSparseTensor,
     self: Tensor,
     min_val: Tensor | int | float,
     max_val: Tensor | int | float,
-) -> DiagonalSparseTensor:
-    if isinstance(self, DiagonalSparseTensor):
+) -> StructuredSparseTensor:
+    if isinstance(self, StructuredSparseTensor):
         raise NotImplementedError()
 
     new_physical = aten.hardtanh_backward.default(grad_output.physical, self, min_val, max_val)
-    return DiagonalSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)
 
 
-@DiagonalSparseTensor.implements(aten.hardswish_backward.default)
-def hardswish_backward_default(grad_output: DiagonalSparseTensor, self: Tensor):
-    if isinstance(self, DiagonalSparseTensor):
+@StructuredSparseTensor.implements(aten.hardswish_backward.default)
+def hardswish_backward_default(grad_output: StructuredSparseTensor, self: Tensor):
+    if isinstance(self, StructuredSparseTensor):
         raise NotImplementedError()
 
     new_physical = aten.hardswish_backward.default(grad_output.physical, self)
-    return DiagonalSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)

src/torchjd/sparse/_aten_function_overrides/einsum.py

@@ -2,23 +2,23 @@
 from torch import Tensor, tensor
 from torch.ops import aten  # type: ignore
 
-from torchjd.sparse import DiagonalSparseTensor
-from torchjd.sparse._diagonal_sparse_tensor import (
+from torchjd.sparse._structured_sparse_tensor import (
+    StructuredSparseTensor,
     p_to_vs_from_v_to_ps,
-    to_diagonal_sparse_tensor,
     to_most_efficient_tensor,
+    to_structured_sparse_tensor,
 )
 
 
 def prepare_for_elementwise_op(
     t1: Tensor | int | float, t2: Tensor | int | float
-) -> tuple[DiagonalSparseTensor, DiagonalSparseTensor]:
+) -> tuple[StructuredSparseTensor, StructuredSparseTensor]:
     """
-    Prepares two DSTs of the same shape from two args, one of those being a DST, and the other being
-    a DST, Tensor, int or float.
+    Prepares two SSTs of the same shape from two args, one of those being a SST, and the other being
+    a SST, Tensor, int or float.
     """
 
-    assert isinstance(t1, DiagonalSparseTensor) or isinstance(t2, DiagonalSparseTensor)
+    assert isinstance(t1, StructuredSparseTensor) or isinstance(t2, StructuredSparseTensor)
 
     if isinstance(t1, int) or isinstance(t1, float):
         t1_ = tensor(t1, device=t2.device)
@@ -31,52 +31,52 @@ def prepare_for_elementwise_op(
         t2_ = t2
 
     t1_, t2_ = aten.broadcast_tensors.default([t1_, t2_])
-    t1_ = to_diagonal_sparse_tensor(t1_)
-    t2_ = to_diagonal_sparse_tensor(t2_)
+    t1_ = to_structured_sparse_tensor(t1_)
+    t2_ = to_structured_sparse_tensor(t2_)
 
     return t1_, t2_
 
 
-@DiagonalSparseTensor.implements(aten.mul.Tensor)
+@StructuredSparseTensor.implements(aten.mul.Tensor)
 def mul_Tensor(t1: Tensor | int | float, t2: Tensor | int | float) -> Tensor:
     # Element-wise multiplication with broadcasting
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
     all_dims = list(range(t1_.ndim))
     return einsum((t1_, all_dims), (t2_, all_dims), output=all_dims)
 
 
-@DiagonalSparseTensor.implements(aten.div.Tensor)
+@StructuredSparseTensor.implements(aten.div.Tensor)
 def div_Tensor(t1: Tensor | int | float, t2: Tensor | int | float) -> Tensor:
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
-    t2_ = DiagonalSparseTensor(1.0 / t2_.physical, t2_.v_to_ps)
+    t2_ = StructuredSparseTensor(1.0 / t2_.physical, t2_.v_to_ps)
     all_dims = list(range(t1_.ndim))
     return einsum((t1_, all_dims), (t2_, all_dims), output=all_dims)
 
 
-@DiagonalSparseTensor.implements(aten.mul.Scalar)
-def mul_Scalar(t: DiagonalSparseTensor, scalar) -> DiagonalSparseTensor:
-    # TODO: maybe it could be that scalar is a scalar DST and t is a normal tensor. Need to check
+@StructuredSparseTensor.implements(aten.mul.Scalar)
+def mul_Scalar(t: StructuredSparseTensor, scalar) -> StructuredSparseTensor:
+    # TODO: maybe it could be that scalar is a scalar SST and t is a normal tensor. Need to check
     #  that
 
-    assert isinstance(t, DiagonalSparseTensor)
+    assert isinstance(t, StructuredSparseTensor)
     new_physical = aten.mul.Scalar(t.physical, scalar)
-    return DiagonalSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.v_to_ps)
 
 
-@DiagonalSparseTensor.implements(aten.add.Tensor)
+@StructuredSparseTensor.implements(aten.add.Tensor)
 def add_Tensor(
     t1: Tensor | int | float, t2: Tensor | int | float, alpha: Tensor | float = 1.0
-) -> DiagonalSparseTensor:
+) -> StructuredSparseTensor:
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
 
     if t1_.v_to_ps == t2_.v_to_ps:
         new_physical = t1_.physical + t2_.physical * alpha
-        return DiagonalSparseTensor(new_physical, t1_.v_to_ps)
+        return StructuredSparseTensor(new_physical, t1_.v_to_ps)
     else:
         raise NotImplementedError()
 
 
-def einsum(*args: tuple[DiagonalSparseTensor, list[int]], output: list[int]) -> Tensor:
+def einsum(*args: tuple[StructuredSparseTensor, list[int]], output: list[int]) -> Tensor:
 
     # First part of the algorithm, determine how to cluster physical indices as well as the common
     # p_shapes corresponding to matching v_dims. Second part translates to physical einsum.
@@ -89,7 +89,7 @@ def einsum(*args: tuple[DiagonalSparseTensor, list[int]], output: list[int]) ->
     # get unique indices
     # map output indices (there can be splits)
     # call physical einsum
-    # build resulting dst
+    # build resulting sst
 
     # OVER
 
@@ -104,7 +104,7 @@ def einsum(*args: tuple[DiagonalSparseTensor, list[int]], output: list[int]) ->
     # [p_1, ..., p_k], then we have to create fresh sub-indices for each dimension.
     # For this reason, an index is decomposed into sub-indices that are then independently
     # clustered.
-    # So if an index i in args for some DiagonalSparseTensor corresponds to a v_to_ps [j, k, l],
+    # So if an index i in args for some StructuredSparseTensor corresponds to a v_to_ps [j, k, l],
     # We will consider three indices (i, 0), (i, 1) and (i, 2).
     # If furthermore [k] correspond to the v_to_ps of some other tensor with index j, then
     # (i, 1) and (j, 0) will be clustered together (and end up being mapped to the same indice in
@@ -136,7 +136,7 @@ def group_indices(indices: list[tuple[int, int]]) -> None:
     tensors = list[Tensor]()
     indices_to_n_pdims = dict[int, int]()
     for t, indices in args:
-        assert isinstance(t, DiagonalSparseTensor)
+        assert isinstance(t, StructuredSparseTensor)
         tensors.append(t.physical)
         for ps, index in zip(t.v_to_ps, indices):
             if index in indices_to_n_pdims:
@@ -150,7 +150,7 @@ def group_indices(indices: list[tuple[int, int]]) -> None:
             group_indices([(indices[i], sub_i) for i, sub_i in indices_])
         # record the physical dimensions, index[v] for v in vs will end-up mapping to the same
         # final dimension as they were just clustered, so we can take the first, which exists as
-        # t is a valid DST.
+        # t is a valid SST.
         new_indices_pair.append([(indices[vs[0][0]], vs[0][1]) for vs in p_to_vs])
 
     current = 0
@@ -186,52 +186,52 @@ def unique_int(pair: tuple[int, int]) -> int:
     return to_most_efficient_tensor(physical, v_to_ps)
 
 
-@DiagonalSparseTensor.implements(aten.bmm.default)
+@StructuredSparseTensor.implements(aten.bmm.default)
 def bmm_default(mat1: Tensor, mat2: Tensor) -> Tensor:
-    assert isinstance(mat1, DiagonalSparseTensor) or isinstance(mat2, DiagonalSparseTensor)
+    assert isinstance(mat1, StructuredSparseTensor) or isinstance(mat2, StructuredSparseTensor)
     assert (
         mat1.ndim == 3
         and mat2.ndim == 3
         and mat1.shape[0] == mat2.shape[0]
         and mat1.shape[2] == mat2.shape[1]
     )
 
-    mat1_ = to_diagonal_sparse_tensor(mat1)
-    mat2_ = to_diagonal_sparse_tensor(mat2)
+    mat1_ = to_structured_sparse_tensor(mat1)
+    mat2_ = to_structured_sparse_tensor(mat2)
 
     # TODO: Verify that the dimension `0` of mat1_ and mat2_ have the same physical dimension sizes
     #  decompositions. If not, can reshape to common decomposition?
     return einsum((mat1_, [0, 1, 2]), (mat2_, [0, 2, 3]), output=[0, 1, 3])
 
 
-@DiagonalSparseTensor.implements(aten.mm.default)
+@StructuredSparseTensor.implements(aten.mm.default)
 def mm_default(mat1: Tensor, mat2: Tensor) -> Tensor:
-    assert isinstance(mat1, DiagonalSparseTensor) or isinstance(mat2, DiagonalSparseTensor)
+    assert isinstance(mat1, StructuredSparseTensor) or isinstance(mat2, StructuredSparseTensor)
     assert mat1.ndim == 2 and mat2.ndim == 2 and mat1.shape[1] == mat2.shape[0]
 
-    mat1_ = to_diagonal_sparse_tensor(mat1)
-    mat2_ = to_diagonal_sparse_tensor(mat2)
+    mat1_ = to_structured_sparse_tensor(mat1)
+    mat2_ = to_structured_sparse_tensor(mat2)
 
     return einsum((mat1_, [0, 1]), (mat2_, [1, 2]), output=[0, 2])
 
 
-@DiagonalSparseTensor.implements(aten.mean.default)
-def mean_default(t: DiagonalSparseTensor) -> Tensor:
-    assert isinstance(t, DiagonalSparseTensor)
+@StructuredSparseTensor.implements(aten.mean.default)
+def mean_default(t: StructuredSparseTensor) -> Tensor:
+    assert isinstance(t, StructuredSparseTensor)
     return aten.sum.default(t.physical) / t.numel()
 
 
-@DiagonalSparseTensor.implements(aten.sum.default)
-def sum_default(t: DiagonalSparseTensor) -> Tensor:
-    assert isinstance(t, DiagonalSparseTensor)
+@StructuredSparseTensor.implements(aten.sum.default)
+def sum_default(t: StructuredSparseTensor) -> Tensor:
+    assert isinstance(t, StructuredSparseTensor)
     return aten.sum.default(t.physical)
 
 
-@DiagonalSparseTensor.implements(aten.sum.dim_IntList)
+@StructuredSparseTensor.implements(aten.sum.dim_IntList)
 def sum_dim_IntList(
-    t: DiagonalSparseTensor, dim: list[int], keepdim: bool = False, dtype=None
+    t: StructuredSparseTensor, dim: list[int], keepdim: bool = False, dtype=None
 ) -> Tensor:
-    assert isinstance(t, DiagonalSparseTensor)
+    assert isinstance(t, StructuredSparseTensor)
 
     if dtype:
         raise NotImplementedError()

src/torchjd/sparse/_aten_function_overrides/pointwise.py

@@ -1,6 +1,6 @@
 from torch.ops import aten  # type: ignore
 
-from torchjd.sparse import DiagonalSparseTensor
+from torchjd.sparse._structured_sparse_tensor import StructuredSparseTensor
 
 # pointwise functions applied to one Tensor with `0.0 → 0`
 _POINTWISE_FUNCTIONS = [
@@ -68,18 +68,18 @@
 
 
 def _override_pointwise(op):
-    @DiagonalSparseTensor.implements(op)
-    def func_(t: DiagonalSparseTensor) -> DiagonalSparseTensor:
-        assert isinstance(t, DiagonalSparseTensor)
-        return DiagonalSparseTensor(op(t.physical), t.v_to_ps)
+    @StructuredSparseTensor.implements(op)
+    def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
+        assert isinstance(t, StructuredSparseTensor)
+        return StructuredSparseTensor(op(t.physical), t.v_to_ps)
 
     return func_
 
 
 def _override_inplace_pointwise(op):
-    @DiagonalSparseTensor.implements(op)
-    def func_(t: DiagonalSparseTensor) -> DiagonalSparseTensor:
-        assert isinstance(t, DiagonalSparseTensor)
+    @StructuredSparseTensor.implements(op)
+    def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
+        assert isinstance(t, StructuredSparseTensor)
         op(t.physical)
         return t
 
@@ -91,22 +91,22 @@ def func_(t: DiagonalSparseTensor) -> DiagonalSparseTensor:
     _override_inplace_pointwise(pointwise_func)
 
 
-@DiagonalSparseTensor.implements(aten.pow.Tensor_Scalar)
-def pow_Tensor_Scalar(t: DiagonalSparseTensor, exponent: float) -> DiagonalSparseTensor:
-    assert isinstance(t, DiagonalSparseTensor)
+@StructuredSparseTensor.implements(aten.pow.Tensor_Scalar)
+def pow_Tensor_Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredSparseTensor:
+    assert isinstance(t, StructuredSparseTensor)
 
     if exponent <= 0.0:
         # Need to densify because we don't have pow(0.0, exponent) = 0.0
         return aten.pow.Tensor_Scalar(t.to_dense(), exponent)
 
     new_physical = aten.pow.Tensor_Scalar(t.physical, exponent)
-    return DiagonalSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.v_to_ps)
 
 
 # Somehow there's no pow_.Tensor_Scalar and pow_.Scalar takes tensor and scalar.
-@DiagonalSparseTensor.implements(aten.pow_.Scalar)
-def pow__Scalar(t: DiagonalSparseTensor, exponent: float) -> DiagonalSparseTensor:
-    assert isinstance(t, DiagonalSparseTensor)
+@StructuredSparseTensor.implements(aten.pow_.Scalar)
+def pow__Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredSparseTensor:
+    assert isinstance(t, StructuredSparseTensor)
 
     if exponent <= 0.0:
         # Need to densify because we don't have pow(0.0, exponent) = 0.0
@@ -117,9 +117,9 @@ def pow__Scalar(t: DiagonalSparseTensor, exponent: float) -> DiagonalSparseTenso
     return t
 
 
-@DiagonalSparseTensor.implements(aten.div.Scalar)
-def div_Scalar(t: DiagonalSparseTensor, divisor: float) -> DiagonalSparseTensor:
-    assert isinstance(t, DiagonalSparseTensor)
+@StructuredSparseTensor.implements(aten.div.Scalar)
+def div_Scalar(t: StructuredSparseTensor, divisor: float) -> StructuredSparseTensor:
+    assert isinstance(t, StructuredSparseTensor)
 
     new_physical = aten.div.Scalar(t.physical, divisor)
-    return DiagonalSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.v_to_ps)