update interface of unsquash_pdim

Author: mattbuot

src/torchjd/sparse/_aten_function_overrides/shape.py

@@ -69,34 +69,59 @@ def infer_shape(shape: list[int], numel: int) -> list[int]:
     return [inferred if s == -1 else s for s in shape]
 
 
-def unsquash_pdim_from_strides(
-    physical: Tensor, pdim: int, new_pdim_shape: list[int]
+def unsquash_pdim(
+    physical: Tensor, strides: Tensor, pdim: int, new_pdim_shape: list[int]
 ) -> tuple[Tensor, Tensor]:
-    new_shape = list(physical.shape)
-    new_shape = new_shape[:pdim] + new_pdim_shape + new_shape[pdim + 1 :]
-    new_physical = physical.reshape(new_shape)
-
-    stride_multipliers = tensor([prod(new_pdim_shape[i + 1 :]) for i in range(len(new_pdim_shape))])
-    return new_physical, stride_multipliers
+    """
+    EXAMPLE:
+
+    physical = [
+        [1, 2, 3, 4, 5, 6],
+        [7, 8, 9, 10, 11, 12],
+        [13, 14, 15, 16, 17, 18],
+    ]
+    strides = [
+        [1, 1],
+        [0, 2],
+    ]
+
+    dim = 1
+    shape = [2, 3]
+
+    new_physical = [[
+        [1, 2, 3],
+        [4, 5, 6],
+    ], [
+        [7, 8, 9],
+        [10, 11, 12],
+    ], [
+        [13, 14, 15],
+        [16, 17, 18],
+    ]]
+
+    new_strides = [
+        [1, 3, 1],
+        [0, 6, 2]
+    """
 
+    # TODO: handle working with multiple dimensions at once
 
-def unsquash_pdim(
-    physical: Tensor, pdim: int, new_pdim_shape: list[int]
-) -> tuple[Tensor, list[list[int]]]:
-    new_shape = list(physical.shape)
-    new_shape = new_shape[:pdim] + new_pdim_shape + new_shape[pdim + 1 :]
+    old_shape = list(physical.shape)
+    new_shape = old_shape[:pdim] + new_pdim_shape + old_shape[pdim + 1 :]
     new_physical = physical.reshape(new_shape)
 
-    def new_encoding_fn(d: int) -> list[int]:
-        if d < pdim:
-            return [d]
-        elif d > pdim:
-            return [d + len(new_pdim_shape) - 1]
-        else:
-            return [pdim + i for i in range(len(new_pdim_shape))]
+    stride_multipliers = tensor([prod(new_pdim_shape[i + 1 :]) for i in range(len(new_pdim_shape))])
+
+    new_strides = torch.concat(
+        [
+            strides[:, :pdim],
+            torch.outer(strides[:, pdim], stride_multipliers),
+            strides[:, pdim + 1 :],
+        ],
+        dim=1,
+    )
 
-    new_encoding = [new_encoding_fn(d) for d in range(len(physical.shape))]
-    return new_physical, new_encoding
+    return new_physical, new_strides
 
 
 @impl(aten._unsafe_view.default)

tests/unit/sparse/test_structured_sparse_tensor.py

@@ -315,29 +315,38 @@ def test_fix_ungrouped_dims(
 @mark.parametrize(
     [
         "physical_shape",
+        "strides",
         "pdim",
         "new_pdim_shape",
         "expected_physical_shape",
-        "expected_new_encoding",
+        "expected_strides",
     ],
     [
-        ([4], 0, [4], [4], [[0]]),  # trivial
-        ([4], 0, [2, 2], [2, 2], [[0, 1]]),
-        ([3, 4, 5], 1, [2, 1, 1, 2], [3, 2, 1, 1, 2, 5], [[0], [1, 2, 3, 4], [5]]),
+        ([4], tensor([[1], [2]]), 0, [4], [4], tensor([[1], [2]])),  # trivial
+        ([4], tensor([[1], [2]]), 0, [2, 2], [2, 2], tensor([[2, 1], [4, 2]])),
+        (
+            [3, 4, 5],
+            tensor([[1, 2, 0], [1, 0, 1], [0, 1, 1]]),
+            1,
+            [2, 1, 1, 2],
+            [3, 2, 1, 1, 2, 5],
+            tensor([[1, 4, 4, 4, 2, 0], [1, 0, 0, 0, 0, 1], [0, 2, 2, 2, 1, 1]]),
+        ),
     ],
 )
 def test_unsquash_pdim(
     physical_shape: list[int],
+    strides: Tensor,
     pdim: int,
     new_pdim_shape: list[int],
     expected_physical_shape: list[int],
-    expected_new_encoding: list[list[int]],
+    expected_strides: Tensor,
 ):
     physical = randn_(physical_shape)
-    new_physical, new_encoding = unsquash_pdim(physical, pdim, new_pdim_shape)
+    new_physical, new_strides = unsquash_pdim(physical, strides, pdim, new_pdim_shape)
 
     assert list(new_physical.shape) == expected_physical_shape
-    assert new_encoding == expected_new_encoding
+    assert torch.equal(new_strides, expected_strides)
 
 
 @mark.parametrize(