Move fix_zero_stride_columns to _coalesce.py

Author: ValerianRey

src/torchjd/sparse/_coalesce.py

@@ -0,0 +1,19 @@
+import torch
+from torch import Tensor
+
+
+def fix_zero_stride_columns(physical: Tensor, strides: Tensor) -> tuple[Tensor, Tensor]:
+    """
+    Remove columns of strides that are all 0 and sum the corresponding elements in the physical
+    tensor.
+    """
+
+    are_columns_zero = (strides == 0).all(dim=0)
+
+    if not are_columns_zero.any():
+        return physical, strides
+
+    zero_column_indices = torch.arange(len(are_columns_zero))[are_columns_zero].tolist()
+    physical = physical.sum(dim=zero_column_indices)
+    strides = strides[:, ~are_columns_zero]
+    return physical, strides

src/torchjd/sparse/_structured_sparse_tensor.py

@@ -276,20 +276,3 @@ def make_sst(physical: Tensor, strides: Tensor) -> StructuredSparseTensor:
     physical, strides = fix_dim_of_size_1(physical, strides)
     physical, strides = fix_ungrouped_dims(physical, strides)
     return StructuredSparseTensor(physical, strides)
-
-
-def fix_zero_stride_columns(physical: Tensor, strides: Tensor) -> tuple[Tensor, Tensor]:
-    """
-    Remove columns of strides that are all 0 and sum the corresponding elements in the physical
-    tensor.
-    """
-
-    are_columns_zero = (strides == 0).all(dim=0)
-
-    if not are_columns_zero.any():
-        return physical, strides
-
-    zero_column_indices = torch.arange(len(are_columns_zero))[are_columns_zero].tolist()
-    physical = physical.sum(dim=zero_column_indices)
-    strides = strides[:, ~are_columns_zero]
-    return physical, strides

tests/unit/sparse/test_structured_sparse_tensor.py

@@ -11,10 +11,10 @@
     _POINTWISE_FUNCTIONS,
 )
 from torchjd.sparse._aten_function_overrides.shape import unsquash_pdim
+from torchjd.sparse._coalesce import fix_zero_stride_columns
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
     fix_ungrouped_dims,
-    fix_zero_stride_columns,
     get_full_source,
     get_groupings,
 )