Fix to_most_efficient_tensor

Author: ValerianRey

src/torchjd/sparse/_structured_sparse_tensor.py

@@ -216,24 +216,8 @@ def to_most_efficient_tensor(physical: Tensor, strides: Tensor) -> Tensor:
     physical, strides = fix_ungrouped_dims(physical, strides)
 
     if (strides.sum(dim=0) == 1).all():
-        # All physical dimensions make you move by 1 in exactly 1 virtual dimension.
-        # Also, because all physical dimensions have been maximally grouped, we cannot have two
-        # physical dimensions that make you move in the same virtual dimension.
-        # So strides is an identity matrix with potentially some extra rows of zeros, and
-        # potentially shuffled columns.
-
-        # The first step is to unsqueeze the physical tensor for each extra row of zeros in the
-        # strides.
-        zero_row_mask = strides.sum(dim=1) == 0
-        number_of_zero_rows = zero_row_mask.sum()
-        for _ in number_of_zero_rows:
-            physical = physical.unsqueeze(-1)
-
-        # The second step is to re-order the physical dimensions so that the corresponding
-        # strides matrix would be an identity.
-        source = arange(strides.shape[0])
-        destination = strides[zero_row_mask] @ source
-        return physical.movedim(list(source), list(destination))
+        # TODO: this can be done more efficiently (without even creating the SST)
+        return StructuredSparseTensor(physical, strides).to_dense()
     else:
         return StructuredSparseTensor(physical, strides)