Add intdiv_c

Author: ValerianRey

src/torchjd/sparse/_linalg.py

@@ -27,7 +27,7 @@ def solve_int(A: Tensor, B: Tensor, tol=1e-9) -> Tensor | None:
 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
+    t1 = d * t2 + r with d = t1 //c t2 and
     0 <= r[i] <= t1[i] for all i.
 
     :param t1: Non-negative integer vector.
@@ -48,3 +48,29 @@ def mod_c(t1: Tensor, t2: Tensor) -> Tensor:
     else:
         min_divider = (t1[non_zero_indices] // t2[non_zero_indices]).min()
         return t1 - min_divider * 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 [0, 0]^T => ValueError
+        [8, 12]^T //c [2, 0]^T = 4
+        [8, 12]^T //c [0, 2]^T = 6
+    """
+
+    non_zero_indices = torch.nonzero(t2)
+    if len(non_zero_indices) == 0:
+        raise ValueError("Cannot divide by the zero vector.")
+    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,7 +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._linalg import intdiv_c, mod_c
 from torchjd.sparse._structured_sparse_tensor import (
     StructuredSparseTensor,
     fix_ungrouped_dims,
@@ -440,3 +440,26 @@ def test_mod_c(
     expected: Tensor,
 ):
     assert torch.equal(mod_c(t1, t2), expected)
+
+
+@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(ValueError):
+        intdiv_c(tensor([3, 4]), tensor([0, 0]))