Add intdiv_c and mod_c

* Note that in python, x % 0 raises ZeroDivisionError. The
  implementation of mod_c matches this behavior when t2 is the zero
vector.

Author: ValerianRey

src/torchjd/sparse/_linalg.py

@@ -22,3 +22,50 @@ 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 = d * t2 + r with d = t1 //c t2 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 [2, 0]^T = [0, 12]^T
+        [8, 12]^T %c [0, 2]^T = [8, 0]^T
+        [8, 12]^T %c [0, 0]^T => ZeroDivisionError
+    """
+
+    return t1 - intdiv_c(t1, t2) * t2
+
+
+def intdiv_c(t1: Tensor, t2: Tensor) -> Tensor:
+    """
+    Computes the combined integer division d = t1 // t2, such that
+    t1 = d * t2 + r with r = t1 %c t2
+    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 = 4
+        [8, 12]^T //c [2, 4]^T = 3
+        [8, 12]^T //c [3, 3]^T = 2
+        [8, 12]^T //c [2, 0]^T = 4
+        [8, 12]^T //c [0, 2]^T = 6
+        [8, 12]^T //c [0, 0]^T => ZeroDivisionError
+    """
+
+    non_zero_indices = torch.nonzero(t2)
+    if len(non_zero_indices) == 0:
+        raise ZeroDivisionError("division by zero")
+    else:
+        min_divider = (t1[non_zero_indices] // t2[non_zero_indices]).min()
+        return min_divider

tests/unit/sparse/test_structured_sparse_tensor.py

@@ -1,5 +1,5 @@
 import torch
-from pytest import mark
+from pytest import mark, raises
 from torch import Tensor, tensor
 from torch.ops import aten  # type: ignore
 from torch.testing import assert_close
@@ -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 intdiv_c, mod_c
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
     fix_ungrouped_dims,
@@ -420,3 +421,49 @@ 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([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)
+
+
+def test_mod_c_by_0_raises():
+    with raises(ZeroDivisionError):
+        mod_c(tensor([3, 4]), tensor([0, 0]))
+
+
+@mark.parametrize(
+    ["t1", "t2", "expected"],
+    [
+        (tensor([8, 12]), tensor([2, 3]), 4),
+        (tensor([8, 12]), tensor([2, 4]), 3),
+        (tensor([8, 12]), tensor([3, 3]), 2),
+        (tensor([8, 12]), tensor([2, 0]), 4),
+        (tensor([8, 12]), tensor([0, 2]), 6),
+    ],
+)
+def test_intdiv_c(
+    t1: Tensor,
+    t2: Tensor,
+    expected: Tensor,
+):
+    assert intdiv_c(t1, t2) == expected
+
+
+def test_intdiv_c_by_0_raises():
+    with raises(ZeroDivisionError):
+        intdiv_c(tensor([3, 4]), tensor([0, 0]))