Fix concat for cases where it has to densify

Author: ValerianRey

src/torchjd/sparse/_aten_function_overrides/shape.py

@@ -127,6 +127,17 @@ def cat_default(tensors: list[Tensor], dim: int = 0) -> Tensor:
             f"basis. Found the following tensors:\n{[repr(t) for t in tensors_]} and the following "
             f"dim: {dim}."
         )
+    if any(t.physical.shape != ref_tensor.physical.shape for t in tensors_[1:]):
+        # This can happen in the following example:
+        # t1 = SLT([1 2 3], [[2]])
+        # t2 = SLT([4 5 6 7], [[2]])
+        # The expected result would be 1 0 2 0 3 4 0 5 0 6 0 7, but this is not representable
+        # efficiently as an SLT (because there is no 0 between 3 and 4, and both physicals have a
+        # different shape so we can't just stack them).
+
+        # TODO: Maybe a partial densify is possible rather than a full densify.
+        print_fallback(aten.cat.default, (tensors, dim), {})
+        return aten.cat.default([unwrap_to_dense(t) for t in tensors])
 
     # We need to try to find the (pretty sure it either does not exist or is unique) physical
     # dimension that makes us only move on virtual dimension dim. It also needs to be such that

tests/unit/sparse/test_sparse_latticed_tensor.py

@@ -332,21 +332,27 @@ def test_get_column_indices(source: list[int], destination: list[int], ndim: int
 
 
 @mark.parametrize(
-    ["slt_args", "dim"],
+    ["slt_args", "dim", "expected_densify"],
     [
-        ([([3], tensor([[1], [1]])), ([3], tensor([[1], [1]]))], 1),
-        ([([3, 2], tensor([[1, 0], [1, 3]])), ([3, 2], tensor([[1, 0], [1, 3]]))], 1),
+        ([([3], tensor([[1], [1]])), ([3], tensor([[1], [1]]))], 1, False),
+        ([([3], tensor([[2]])), ([4], tensor([[2]]))], 0, True),
+        ([([3, 2], tensor([[1, 0], [1, 3]])), ([3, 2], tensor([[1, 0], [1, 3]]))], 1, False),
     ],
 )
 def test_concatenate(
     slt_args: list[tuple[list[int], Tensor]],
     dim: int,
+    expected_densify: bool,
 ):
     tensors = [SparseLatticedTensor(randn_(pshape), basis) for pshape, basis in slt_args]
     res = aten.cat.default(tensors, dim)
     expected = aten.cat.default([t.to_dense() for t in tensors], dim)
 
-    assert isinstance(res, SparseLatticedTensor)
+    if expected_densify:
+        assert not isinstance(res, SparseLatticedTensor)
+    else:
+        assert isinstance(res, SparseLatticedTensor)
+
     assert torch.all(torch.eq(res.to_dense(), expected))