Use strides-based representation instead of v_to_ps-based

* still need to update einsum and view, update tests, debug, and maybe delete some unused functions.

Author: ValerianRey

src/torchjd/sparse/_aten_function_overrides/backward.py

@@ -10,7 +10,7 @@ def threshold_backward_default(
 ) -> StructuredSparseTensor:
     new_physical = aten.threshold_backward.default(grad_output.physical, self, threshold)
 
-    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.strides)
 
 
 @impl(aten.hardtanh_backward.default)
@@ -24,7 +24,7 @@ def hardtanh_backward_default(
         raise NotImplementedError()
 
     new_physical = aten.hardtanh_backward.default(grad_output.physical, self, min_val, max_val)
-    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.strides)
 
 
 @impl(aten.hardswish_backward.default)
@@ -33,4 +33,4 @@ def hardswish_backward_default(grad_output: StructuredSparseTensor, self: Tensor
         raise NotImplementedError()
 
     new_physical = aten.hardswish_backward.default(grad_output.physical, self)
-    return StructuredSparseTensor(new_physical, grad_output.v_to_ps)
+    return StructuredSparseTensor(new_physical, grad_output.strides)

src/torchjd/sparse/_aten_function_overrides/einsum.py

@@ -5,7 +5,6 @@
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
     impl,
-    p_to_vs_from_v_to_ps,
     to_most_efficient_tensor,
     to_structured_sparse_tensor,
 )
@@ -84,7 +83,7 @@ def group_indices(indices: list[tuple[int, int]]) -> None:
                     )
             else:
                 indices_to_n_pdims[index] = len(pdims)
-        p_to_vs = p_to_vs_from_v_to_ps(t.v_to_ps)
+        p_to_vs = ...  # p_to_vs_from_v_to_ps(t.v_to_ps)
         for indices_ in p_to_vs:
             # elements in indices[indices_] map to the same dimension, they should be clustered
             # together
@@ -165,7 +164,7 @@ 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_.v_to_ps)
+    t2_ = StructuredSparseTensor(1.0 / t2_.physical, t2_.strides)
     all_dims = list(range(t1_.ndim))
     return einsum((t1_, all_dims), (t2_, all_dims), output=all_dims)
 
@@ -177,7 +176,7 @@ def mul_Scalar(t: StructuredSparseTensor, scalar) -> StructuredSparseTensor:
 
     assert isinstance(t, StructuredSparseTensor)
     new_physical = aten.mul.Scalar(t.physical, scalar)
-    return StructuredSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.strides)
 
 
 @impl(aten.add.Tensor)
@@ -186,9 +185,9 @@ def add_Tensor(
 ) -> StructuredSparseTensor:
     t1_, t2_ = prepare_for_elementwise_op(t1, t2)
 
-    if t1_.v_to_ps == t2_.v_to_ps:
+    if torch.equal(t1_.strides, t2_.strides):
         new_physical = t1_.physical + t2_.physical * alpha
-        return StructuredSparseTensor(new_physical, t1_.v_to_ps)
+        return StructuredSparseTensor(new_physical, t1_.strides)
     else:
         raise NotImplementedError()
 

src/torchjd/sparse/_aten_function_overrides/pointwise.py

@@ -71,7 +71,7 @@ def _override_pointwise(op):
     @impl(op)
     def func_(t: StructuredSparseTensor) -> StructuredSparseTensor:
         assert isinstance(t, StructuredSparseTensor)
-        return StructuredSparseTensor(op(t.physical), t.v_to_ps)
+        return StructuredSparseTensor(op(t.physical), t.strides)
 
     return func_
 
@@ -100,7 +100,7 @@ def pow_Tensor_Scalar(t: StructuredSparseTensor, exponent: float) -> StructuredS
         return aten.pow.Tensor_Scalar(t.to_dense(), exponent)
 
     new_physical = aten.pow.Tensor_Scalar(t.physical, exponent)
-    return StructuredSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.strides)
 
 
 # Somehow there's no pow_.Tensor_Scalar and pow_.Scalar takes tensor and scalar.
@@ -122,4 +122,4 @@ def div_Scalar(t: StructuredSparseTensor, divisor: float) -> StructuredSparseTen
     assert isinstance(t, StructuredSparseTensor)
 
     new_physical = aten.div.Scalar(t.physical, divisor)
-    return StructuredSparseTensor(new_physical, t.v_to_ps)
+    return StructuredSparseTensor(new_physical, t.strides)

src/torchjd/sparse/_aten_function_overrides/shape.py

@@ -2,13 +2,11 @@
 from typing import cast
 
 import torch
-from torch import Tensor, tensor
+from torch import Tensor, arange, tensor
 from torch.ops import aten  # type: ignore
 
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
-    encode_v_to_ps,
-    fix_dim_encoding,
     impl,
     print_fallback,
     to_most_efficient_tensor,
@@ -136,10 +134,10 @@ def unsqueeze_default(t: StructuredSparseTensor, dim: int) -> StructuredSparseTe
     if dim < 0:
         dim = t.ndim + dim + 1
 
-    new_v_to_ps = [p for p in t.v_to_ps]  # Deepcopy the list to not modify the original v_to_ps
-    new_v_to_ps.insert(dim, [])
-
-    return StructuredSparseTensor(t.physical, new_v_to_ps)
+    new_strides = torch.concatenate(
+        [t.strides[:dim], torch.zeros(1, t.strides.shape[1], dtype=torch.int64), t.strides[dim:]]
+    )
+    return StructuredSparseTensor(t.physical, new_strides)
 
 
 @impl(aten.squeeze.dims)
@@ -153,17 +151,15 @@ def squeeze_dims(t: StructuredSparseTensor, dims: list[int] | int | None) -> Ten
     else:
         excluded = set(dims)
 
-    new_v_to_ps = [pdims for i, pdims in enumerate(t.v_to_ps) if i not in excluded]
-
-    return to_most_efficient_tensor(t.physical, new_v_to_ps)
+    is_row_kept = [i not in excluded for i in range(t.ndim)]
+    new_strides = t.strides[is_row_kept]
+    return to_most_efficient_tensor(t.physical, new_strides)
 
 
 @impl(aten.permute.default)
 def permute_default(t: StructuredSparseTensor, dims: list[int]) -> StructuredSparseTensor:
-    new_v_to_ps = [t.v_to_ps[d] for d in dims]
-
-    new_physical, new_v_to_ps = fix_dim_encoding(t.physical, new_v_to_ps)
-    return StructuredSparseTensor(new_physical, new_v_to_ps)
+    new_strides = t.strides[torch.tensor(dims)]
+    return StructuredSparseTensor(t.physical, new_strides)
 
 
 @impl(aten.cat.default)
@@ -197,25 +193,20 @@ def cat_default(tensors: list[Tensor], dim: int) -> Tensor:
         # Add a physical dimension pdim on which we can concatenate the physicals such that this
         # translates into a concatenation of the virtuals on virtual dimension dim.
 
-        # Stride-based representation:
-        # new_stride_column = torch.zeros(ref_tensor.ndim, dtype=torch.int)
-        # new_stride_column[dim] = ref_virtual_dim_size
-
         pdim = ref_tensor.physical.ndim
-        new_v_to_ps = [[d for d in pdims] for pdims in ref_tensor.v_to_ps]
-        new_v_to_ps[dim] = [pdim] + new_v_to_ps[dim]
-        new_v_to_ps, destination = encode_v_to_ps(new_v_to_ps)
-        source = list(range(len(destination)))
-        physicals = [t.physical.unsqueeze(-1).movedim(source, destination) for t in tensors_]
+        physicals = [t.physical.unsqueeze(-1) for t in tensors_]
+        new_stride_column = torch.zeros(ref_tensor.ndim, 1, dtype=torch.int64)
+        new_stride_column[dim, 0] = ref_virtual_dim_size
+        new_strides = torch.concatenate([ref_tensor.strides, new_stride_column], dim=1)
     else:
         # Such a physical dimension already exists. Note that an alternative implementation would be
         # to simply always add the physical dimension, and squash it if it ends up being not needed.
         physicals = [t.physical for t in tensors_]
         pdim = indices[0][0]
-        new_v_to_ps = ref_tensor.v_to_ps
+        new_strides = ref_tensor.strides
 
     new_physical = aten.cat.default(physicals, dim=pdim)
-    return StructuredSparseTensor(new_physical, new_v_to_ps)
+    return StructuredSparseTensor(new_physical, new_strides)
 
 
 @impl(aten.expand.default)
@@ -225,32 +216,32 @@ def expand_default(t: StructuredSparseTensor, sizes: list[int]) -> StructuredSpa
     assert isinstance(sizes, list)
     assert len(sizes) >= t.ndim
 
+    # Add as many dimensions as needed at the beginning of the tensor (as torch.expand works)
     for _ in range(len(sizes) - t.ndim):
         t = t.unsqueeze(0)
 
-    assert len(sizes) == t.ndim
-
+    # Try to expand each dimension to its new size
     new_physical = t.physical
-    new_v_to_ps = t.v_to_ps
-    n_added_physical_dims = 0
-    for dim, (ps, orig_size, new_size) in enumerate(zip(t.v_to_ps, t.shape, sizes, strict=True)):
-        if len(ps) > 0 and orig_size != new_size and new_size != -1:
+    new_strides = t.strides
+    for d, (vstride, orig_size, new_size) in enumerate(zip(t.strides, t.shape, sizes, strict=True)):
+        if vstride.sum() > 0 and orig_size != new_size and new_size != -1:
             raise ValueError(
-                f"Cannot expand dim {dim} of size != 1. Found size {orig_size} and target size "
+                f"Cannot expand dim {d} of size != 1. Found size {orig_size} and target size "
                 f"{new_size}."
             )
 
-        if len(ps) == 0 and new_size != 1 and new_size != -1:
+        if vstride.sum() == 0 and new_size != 1 and new_size != -1:
             # Add a dimension of size new_size at the end of the physical tensor.
             new_physical_shape = list(new_physical.shape) + [new_size]
             new_physical = new_physical.unsqueeze(-1).expand(new_physical_shape)
-            new_v_to_ps[dim] = [t.physical.ndim + n_added_physical_dims]
-            n_added_physical_dims += 1
 
-    new_v_to_ps, destination = encode_v_to_ps(new_v_to_ps)
-    new_physical = new_physical.movedim(list(range(len(destination))), destination)
+            # Make this new physical dimension have a stride of 1 at virtual dimension d and 0 at
+            # every other virtual dimension
+            new_stride_column = torch.zeros(t.ndim, 1, dtype=torch.int64)
+            new_stride_column[d, 0] = 1
+            new_strides = torch.cat([new_strides, new_stride_column], dim=1)
 
-    return StructuredSparseTensor(new_physical, new_v_to_ps)
+    return StructuredSparseTensor(new_physical, new_strides)
 
 
 @impl(aten.broadcast_tensors.default)
@@ -280,49 +271,14 @@ def broadcast_tensors_default(tensors: list[Tensor]) -> tuple[Tensor, Tensor]:
     return aten.expand.default(t1, new_shape), aten.expand.default(t2, new_shape)
 
 
-@impl(aten.slice.Tensor)
-def slice_Tensor(
-    t: StructuredSparseTensor, dim: int, start: int | None, end: int | None, step: int = 1
-) -> StructuredSparseTensor:
-    assert isinstance(t, StructuredSparseTensor)
-
-    physical_dims = t.v_to_ps[dim]
-
-    if len(physical_dims) > 1:
-        raise ValueError(
-            "Cannot yet slice virtual dim corresponding to several physical dims.\n"
-            f"{t.debug_info()}\n"
-            f"dim={dim}, start={start}, end={end}, step={step}."
-        )
-    elif len(physical_dims) == 0:
-        # Trying to slice a virtual dim of size 1.
-        # Either
-        # - the element of this dim is included in the slice: keep it as it is
-        # - it's not included in the slice (e.g. end<=start): we would end up with a size of 0 on
-        #   that dimension, so we'd need to add a dimension of size 0 to the physical. This is not
-        #   implemented yet.
-        start_ = start if start is not None else 0
-        end_ = end if end is not None else 1
-        if end_ <= start_:  # TODO: the condition might be a bit more complex if step != 1
-            raise NotImplementedError(
-                "Slicing of dimension of size 1 leading to dimension of size 0 not implemented yet."
-            )
-        else:
-            new_physical = t.physical
-    else:
-        physical_dim = physical_dims[0]
-        new_physical = aten.slice.Tensor(t.physical, physical_dim, start, end, step)
-
-    return StructuredSparseTensor(new_physical, t.v_to_ps)
-
-
 @impl(aten.transpose.int)
 def transpose_int(t: StructuredSparseTensor, dim0: int, dim1: int) -> StructuredSparseTensor:
     assert isinstance(t, StructuredSparseTensor)
+    return StructuredSparseTensor(t.physical, _swap_rows(t.strides, dim0, dim1))
 
-    new_v_to_ps = [dims for dims in t.v_to_ps]
-    new_v_to_ps[dim0] = t.v_to_ps[dim1]
-    new_v_to_ps[dim1] = t.v_to_ps[dim0]
 
-    new_physical, new_v_to_ps = fix_dim_encoding(t.physical, new_v_to_ps)
-    return StructuredSparseTensor(new_physical, new_v_to_ps)
+def _swap_rows(matrix: Tensor, c0: int, c1: int) -> Tensor:
+    index = arange(matrix.shape[0])
+    index[c0] = c1
+    index[c1] = c0
+    return matrix[index]