Add initial implementation of view_default

Author: ValerianRey

src/torchjd/sparse/_aten_function_overrides/shape.py

@@ -1,3 +1,5 @@
+import operator
+from itertools import accumulate
 from math import prod
 from typing import cast
 
@@ -16,66 +18,53 @@
 
 @impl(aten.view.default)
 def view_default(t: StructuredSparseTensor, shape: list[int]) -> Tensor:
+    """
+    The main condition that we want to respect is that the indexing in the flattened virtual
+    tensor should remain the same before and after the reshape, i.e.
+
+                                    c.T S = c'.T S'    (1)
+    where:
+    * c is the reversed vector of cumulative physical shape before the reshape, i.e.
+        c.T = [prod(t.shape[1:]), prod(t.shape[2:]), ..., t.shape[-1], 1]
+    * c' is the same thing but after the reshape, i.e.
+        c'.T = [prod(shape[1:]), prod(shape[2:]), ..., shape[-1], 1]
+    * S is the original matrix of strides (t.strides)
+    * S' is the matrix of strides after reshaping.
+
+    For u, v in Z^m and c in Z, say that u ≡ v (mod c) if u_i ≡ v_i (mod c) for all i.
+    Note that c'.T S' ≡ S'[-1] (mod shape[-1])
+    So if we set S'[-1] = c.T S % shape[-1], we have c.T S ≡ c'.T S' (mod shape[-1])
+
+    (c'.T S' - S'[-1]) // shape[-1] ≡ S'[-1] (mod shape[-1])
+    ...
+    """
+
     assert isinstance(t, StructuredSparseTensor)
 
     shape = infer_shape(shape, t.numel())
 
     if prod(shape) != t.numel():
         raise ValueError(f"shape '{shape}' is invalid for input of size {t.numel()}")
 
-    new_v_to_ps = []
-    idx = 0
-    flat_v_to_ps = [dim for dims in t.v_to_ps for dim in dims]
-    new_physical = t.physical
-    for s in shape:
-        group = []
-        current_size = 1
-
-        while current_size < s:
-            if idx >= len(flat_v_to_ps):
-                # TODO: I don't think this can happen, need to review and remove if I'm right.
-                raise ValueError()
-
-            pdim = flat_v_to_ps[idx]
-            pdim_size = new_physical.shape[pdim]
-
-            if current_size * pdim_size > s:
-                # Need to split physical dimension
-                if s % current_size != 0:
-                    raise ValueError("Can't split physical dimension")
-
-                new_pdim_first_dim_size = s // current_size
-
-                if pdim_size % new_pdim_first_dim_size != 0:
-                    raise ValueError("Can't split physical dimension")
-
-                new_pdim_shape = [new_pdim_first_dim_size, pdim_size // new_pdim_first_dim_size]
-                new_physical, new_encoding = unsquash_pdim(new_physical, pdim, new_pdim_shape)
-
-                new_v_to_ps = [
-                    [new_d for d in dims for new_d in new_encoding[d]] for dims in new_v_to_ps
-                ]
-                # A bit of a weird trick here. We want to re-encode flat_v_to_ps according to
-                # new_encoding. However, re-encoding elements before idx would potentially change
-                # the length of the list before idx, so idx would not have the right value anymore.
-                # Since we don't need the elements of flat_v_to_ps that are before idx anyway, we
-                # just get rid of them and re-encode flat_v_to_ps[idx:] instead, and reset idx to 0
-                # to say that we're back at the beginning of this new list.
-                flat_v_to_ps = [new_d for d in flat_v_to_ps[idx:] for new_d in new_encoding[d]]
-                idx = 0
-
-            group.append(pdim)
-            current_size *= new_physical.shape[pdim]
-            idx += 1
-
-        new_v_to_ps.append(group)
-
-    if idx != len(flat_v_to_ps):
-        raise ValueError(f"idx != len(flat_v_to_ps). {idx}; {flat_v_to_ps}; {shape}; {t.v_to_ps}")
-
-    # The above code does not handle physical dimension squashing, so the physical is not
-    # necessarily maximally squashed at this point, so we need the safe constructor.
-    return to_most_efficient_tensor(new_physical, new_v_to_ps)
+    S = t.strides
+    vshape = list(t.shape)
+    c = _reverse_cumulative_product(vshape)
+    remaining_cT_S = c @ S
+
+    stride_rows = list[Tensor]()
+    for modulo in shape[::-1]:
+        stride_row = remaining_cT_S % modulo
+        stride_rows.append(stride_row)
+        remaining_cT_S = (remaining_cT_S - stride_row) // modulo
+        # I think we could skip the - stride_row because the floor div will handle it for us, but it
+        # will make code harder to understand.
+
+    new_strides = torch.stack(stride_rows, dim=0)
+    return to_most_efficient_tensor(t.physical, new_strides)
+
+
+def _reverse_cumulative_product(values: list[int]) -> Tensor:
+    return tensor(list(accumulate((values[1:] + [1])[::-1], operator.mul))[::-1])
 
 
 def infer_shape(shape: list[int], numel: int) -> list[int]: