Add mod_c

Author: ValerianRey

src/torchjd/sparse/_linalg.py

@@ -22,3 +22,29 @@ def solve_int(A: Tensor, B: Tensor, tol=1e-9) -> Tensor | None:
 
     # TODO: Verify that the round operation cannot fail
     return X_rounded.to(torch.int64)
+
+
+def mod_c(t1: Tensor, t2: Tensor) -> Tensor:
+    """
+    Computes the combined modulo r = t1 %c t2, such that
+    t1 = k * t2 + r with k any integer and
+    0 <= r[i] <= t1[i] for all i.
+
+    :param t1: Non-negative integer vector.
+    :param t2: Non-negative integer vector.
+
+    Examples:
+        [8, 12]^T %c [2, 3]^T = [0, 0]^T
+        [8, 12]^T %c [2, 4]^T = [2, 0]^T
+        [8, 12]^T %c [3, 3]^T = [2, 6]^T
+        [8, 12]^T %c [0, 0]^T = [8, 12]^T
+        [8, 12]^T %c [2, 0]^T = [0, 12]^T
+        [8, 12]^T %c [0, 2]^T = [8, 0]^T
+    """
+
+    non_zero_indices = torch.nonzero(t2)
+    if len(non_zero_indices) == 0:
+        return t1
+    else:
+        min_divider = (t1[non_zero_indices] // t2[non_zero_indices]).min()
+        return t1 - min_divider * t2

tests/unit/sparse/test_structured_sparse_tensor.py

@@ -12,6 +12,7 @@
 )
 from torchjd.sparse._aten_function_overrides.shape import unsquash_pdim
 from torchjd.sparse._coalesce import fix_zero_stride_columns
+from torchjd.sparse._linalg import mod_c
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
     fix_ungrouped_dims,
@@ -420,3 +421,22 @@ def test_fix_zero_stride_columns(
     physical, strides = fix_zero_stride_columns(physical, strides)
     assert torch.equal(physical, expected_physical)
     assert torch.equal(strides, expected_strides)
+
+
+@mark.parametrize(
+    ["t1", "t2", "expected"],
+    [
+        (tensor([8, 12]), tensor([2, 3]), tensor([0, 0])),
+        (tensor([8, 12]), tensor([2, 4]), tensor([2, 0])),
+        (tensor([8, 12]), tensor([3, 3]), tensor([2, 6])),
+        (tensor([8, 12]), tensor([0, 0]), tensor([8, 12])),
+        (tensor([8, 12]), tensor([2, 0]), tensor([0, 12])),
+        (tensor([8, 12]), tensor([0, 2]), tensor([8, 0])),
+    ],
+)
+def test_mod_c(
+    t1: Tensor,
+    t2: Tensor,
+    expected: Tensor,
+):
+    assert torch.equal(mod_c(t1, t2), expected)