Rename SST to SparseLatticedTensor

Author: ValerianRey

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 ._structured_sparse_tensor import StructuredSparseTensor, make_sst
+from ._sparse_latticed_tensor import SparseLatticedTensor, 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._structured_sparse_tensor import StructuredSparseTensor, impl
+from torchjd.sparse._sparse_latticed_tensor import SparseLatticedTensor, impl
 
 
 @impl(aten.threshold_backward.default)
 def threshold_backward_default(
-    grad_output: StructuredSparseTensor, self: Tensor, threshold
-) -> StructuredSparseTensor:
+    grad_output: SparseLatticedTensor, self: Tensor, threshold
+) -> SparseLatticedTensor:
     new_physical = aten.threshold_backward.default(grad_output.physical, self, threshold)
 
-    return StructuredSparseTensor(new_physical, grad_output.strides)
+    return SparseLatticedTensor(new_physical, grad_output.strides)
 
 
 @impl(aten.hardtanh_backward.default)
 def hardtanh_backward_default(
-    grad_output: StructuredSparseTensor,
+    grad_output: SparseLatticedTensor,
     self: Tensor,
     min_val: Tensor | int | float,
     max_val: Tensor | int | float,
-) -> StructuredSparseTensor:
-    if isinstance(self, StructuredSparseTensor):
+) -> SparseLatticedTensor:
+    if isinstance(self, SparseLatticedTensor):
         raise NotImplementedError()
 
     new_physical = aten.hardtanh_backward.default(grad_output.physical, self, min_val, max_val)
-    return StructuredSparseTensor(new_physical, grad_output.strides)
+    return SparseLatticedTensor(new_physical, grad_output.strides)
 
 
 @impl(aten.hardswish_backward.default)
-def hardswish_backward_default(grad_output: StructuredSparseTensor, self: Tensor):
-    if isinstance(self, StructuredSparseTensor):
+def hardswish_backward_default(grad_output: SparseLatticedTensor, self: Tensor):
+    if isinstance(self, SparseLatticedTensor):
         raise NotImplementedError()
 
     new_physical = aten.hardswish_backward.default(grad_output.physical, self)
-    return StructuredSparseTensor(new_physical, grad_output.strides)
+    return SparseLatticedTensor(new_physical, grad_output.strides)

src/torchjd/sparse/_aten_function_overrides/einsum.py

@@ -2,15 +2,15 @@
 from torch import Tensor, tensor
 from torch.ops import aten  # type: ignore
 
-from torchjd.sparse._structured_sparse_tensor import (
-    StructuredSparseTensor,
+from torchjd.sparse._sparse_latticed_tensor import (
+    SparseLatticedTensor,
     impl,
     to_most_efficient_tensor,
-    to_structured_sparse_tensor,
+    to_sparse_latticed_tensor,
 )
 
 
-def einsum(*args: tuple[StructuredSparseTensor, list[int]], output: list[int]) -> Tensor:
+def einsum(*args: tuple[SparseLatticedTensor, list[int]], output: list[int]) -> Tensor:
     raise NotImplementedError()
 
     # First part of the algorithm, determine how to cluster physical indices as well as the common
@@ -39,13 +39,13 @@ def einsum(*args: tuple[StructuredSparseTensor, 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 StructuredSparseTensor corresponds to a v_to_ps [j, k, l],
+    # So if an index i in args for some SparseLatticedTensor 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
     # the resulting einsum).
     # Note that this is a problem if two virtual dimensions (from possibly different
-    # StructuredSparseTensors) have the same size but not the same decomposition into physical
+    # SparseLatticedTensors) have the same size but not the same decomposition into physical
     # dimension sizes. For now lets leave the responsibility to care about that in the calling
     # functions, if we can factor code later on we will.
 
@@ -71,7 +71,7 @@ def group_indices(indices: list[tuple[int, int]]) -> None:
     physicals = list[Tensor]()
     indices_to_n_pdims = dict[int, int]()
     for t, indices in args:
-        assert isinstance(t, StructuredSparseTensor)
+        assert isinstance(t, SparseLatticedTensor)
         physicals.append(t.physical)
         for pdims, index in zip(t.v_to_ps, indices):
             if index in indices_to_n_pdims:
@@ -129,13 +129,13 @@ def unique_int(pair: tuple[int, int]) -> int:
 
 def prepare_for_elementwise_op(
     t1: Tensor | int | float, t2: Tensor | int | float
-) -> tuple[StructuredSparseTensor, StructuredSparseTensor]:
+) -> tuple[SparseLatticedTensor, SparseLatticedTensor]:
     """
     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, StructuredSparseTensor) or isinstance(t2, StructuredSparseTensor)
+    assert isinstance(t1, SparseLatticedTensor) or isinstance(t2, SparseLatticedTensor)
 
     if isinstance(t1, int) or isinstance(t1, float):
         t1_ = tensor(t1, device=t2.device)
@@ -148,8 +148,8 @@ def prepare_for_elementwise_op(
         t2_ = t2
 
     t1_, t2_ = aten.broadcast_tensors.default([t1_, t2_])
-    t1_ = to_structured_sparse_tensor(t1_)
-    t2_ = to_structured_sparse_tensor(t2_)
+    t1_ = to_sparse_latticed_tensor(t1_)
+    t2_ = to_sparse_latticed_tensor(t2_)
 
     return t1_, t2_
 
@@ -165,46 +165,46 @@ def mul_Tensor(t1: Tensor | int | float, t2: Tensor | int | float) -> Tensor:
 @impl(aten.div.Tensor)
 def div_Tensor(t1: Tensor | int | float, t2: Tensor | int | float) -> Tensor:
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
-    t2_ = StructuredSparseTensor(1.0 / t2_.physical, t2_.strides)
+    t2_ = SparseLatticedTensor(1.0 / t2_.physical, t2_.strides)
     all_dims = list(range(t1_.ndim))
     return einsum((t1_, all_dims), (t2_, all_dims), output=all_dims)
 
 
 @impl(aten.mul.Scalar)
-def mul_Scalar(t: StructuredSparseTensor, scalar) -> StructuredSparseTensor:
+def mul_Scalar(t: SparseLatticedTensor, scalar) -> SparseLatticedTensor:
     # TODO: maybe it could be that scalar is a scalar SST and t is a normal tensor. Need to check
     #  that
 
-    assert isinstance(t, StructuredSparseTensor)
+    assert isinstance(t, SparseLatticedTensor)
     new_physical = aten.mul.Scalar(t.physical, scalar)
-    return StructuredSparseTensor(new_physical, t.strides)
+    return SparseLatticedTensor(new_physical, t.strides)
 
 
 @impl(aten.add.Tensor)
 def add_Tensor(
     t1: Tensor | int | float, t2: Tensor | int | float, alpha: Tensor | float = 1.0
-) -> StructuredSparseTensor:
+) -> SparseLatticedTensor:
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
 
     if torch.equal(t1_.strides, t2_.strides):
         new_physical = t1_.physical + t2_.physical * alpha
-        return StructuredSparseTensor(new_physical, t1_.strides)
+        return SparseLatticedTensor(new_physical, t1_.strides)
     else:
         raise NotImplementedError()
 
 
 @impl(aten.bmm.default)
 def bmm_default(mat1: Tensor, mat2: Tensor) -> Tensor:
-    assert isinstance(mat1, StructuredSparseTensor) or isinstance(mat2, StructuredSparseTensor)
+    assert isinstance(mat1, SparseLatticedTensor) or isinstance(mat2, SparseLatticedTensor)
     assert (
         mat1.ndim == 3
         and mat2.ndim == 3
         and mat1.shape[0] == mat2.shape[0]
         and mat1.shape[2] == mat2.shape[1]
     )
 
-    mat1_ = to_structured_sparse_tensor(mat1)
-    mat2_ = to_structured_sparse_tensor(mat2)
+    mat1_ = to_sparse_latticed_tensor(mat1)
+    mat2_ = to_sparse_latticed_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?
@@ -213,32 +213,32 @@ def bmm_default(mat1: Tensor, mat2: Tensor) -> Tensor:
 
 @impl(aten.mm.default)
 def mm_default(mat1: Tensor, mat2: Tensor) -> Tensor:
-    assert isinstance(mat1, StructuredSparseTensor) or isinstance(mat2, StructuredSparseTensor)
+    assert isinstance(mat1, SparseLatticedTensor) or isinstance(mat2, SparseLatticedTensor)
     assert mat1.ndim == 2 and mat2.ndim == 2 and mat1.shape[1] == mat2.shape[0]
 
-    mat1_ = to_structured_sparse_tensor(mat1)
-    mat2_ = to_structured_sparse_tensor(mat2)
+    mat1_ = to_sparse_latticed_tensor(mat1)
+    mat2_ = to_sparse_latticed_tensor(mat2)
 
     return einsum((mat1_, [0, 1]), (mat2_, [1, 2]), output=[0, 2])
 
 
 @impl(aten.mean.default)
-def mean_default(t: StructuredSparseTensor) -> Tensor:
-    assert isinstance(t, StructuredSparseTensor)
+def mean_default(t: SparseLatticedTensor) -> Tensor:
+    assert isinstance(t, SparseLatticedTensor)
     return aten.sum.default(t.physical) / t.numel()
 
 
 @impl(aten.sum.default)
-def sum_default(t: StructuredSparseTensor) -> Tensor:
-    assert isinstance(t, StructuredSparseTensor)
+def sum_default(t: SparseLatticedTensor) -> Tensor:
+    assert isinstance(t, SparseLatticedTensor)
     return aten.sum.default(t.physical)
 
 
 @impl(aten.sum.dim_IntList)
 def sum_dim_IntList(
-    t: StructuredSparseTensor, dim: list[int], keepdim: bool = False, dtype=None
+    t: SparseLatticedTensor, dim: list[int], keepdim: bool = False, dtype=None
 ) -> Tensor:
-    assert isinstance(t, StructuredSparseTensor)
+    assert isinstance(t, SparseLatticedTensor)
 
     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._structured_sparse_tensor import StructuredSparseTensor, impl
+from torchjd.sparse._sparse_latticed_tensor import SparseLatticedTensor, impl
 
 # pointwise functions applied to one Tensor with `0.0 → 0`
 _POINTWISE_FUNCTIONS = [
@@ -69,17 +69,17 @@
 
 def _override_pointwise(op):
     @impl(op)
-    def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
-        assert isinstance(t, StructuredSparseTensor)
-        return StructuredSparseTensor(op(t.physical), t.strides)
+    def func_(t: SparseLatticedTensor) -> SparseLatticedTensor:
+        assert isinstance(t, SparseLatticedTensor)
+        return SparseLatticedTensor(op(t.physical), t.strides)
 
     return func_
 
 
 def _override_inplace_pointwise(op):
     @impl(op)
-    def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
-        assert isinstance(t, StructuredSparseTensor)
+    def func_(t: SparseLatticedTensor) -> SparseLatticedTensor:
+        assert isinstance(t, SparseLatticedTensor)
         op(t.physical)
         return t
 
@@ -92,21 +92,21 @@ def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
 
 
 @impl(aten.pow.Tensor_Scalar)
-def pow_Tensor_Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredSparseTensor:
-    assert isinstance(t, StructuredSparseTensor)
+def pow_Tensor_Scalar(t: SparseLatticedTensor, exponent: float) -> SparseLatticedTensor:
+    assert isinstance(t, SparseLatticedTensor)
 
     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 StructuredSparseTensor(new_physical, t.strides)
+    return SparseLatticedTensor(new_physical, t.strides)
 
 
 # Somehow there's no pow_.Tensor_Scalar and pow_.Scalar takes tensor and scalar.
 @impl(aten.pow_.Scalar)
-def pow__Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredSparseTensor:
-    assert isinstance(t, StructuredSparseTensor)
+def pow__Scalar(t: SparseLatticedTensor, exponent: float) -> SparseLatticedTensor:
+    assert isinstance(t, SparseLatticedTensor)
 
     if exponent <= 0.0:
         # Need to densify because we don't have pow(0.0, exponent) = 0.0
@@ -118,8 +118,8 @@ def pow__Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredSparseT
 
 
 @impl(aten.div.Scalar)
-def div_Scalar(t: StructuredSparseTensor, divisor: float) -> StructuredSparseTensor:
-    assert isinstance(t, StructuredSparseTensor)
+def div_Scalar(t: SparseLatticedTensor, divisor: float) -> SparseLatticedTensor:
+    assert isinstance(t, SparseLatticedTensor)
 
     new_physical = aten.div.Scalar(t.physical, divisor)
-    return StructuredSparseTensor(new_physical, t.strides)
+    return SparseLatticedTensor(new_physical, t.strides)