chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,47 @@
+import backend as F
+import pytest
+import torch
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("dim", [0, 1])
+@pytest.mark.parametrize("index", [None, (1, 3), (4, 0, 2)])
+def test_compact(create_func, dim, index):
+    ctx = F.ctx()
+    shape = (5, 5)
+    ans_idx = []
+    if index is not None:
+        ans_idx = list(dict.fromkeys(index))
+        index = torch.tensor(index).to(ctx)
+
+    A = create_func(shape, 8, ctx)
+
+    A_compact, ret_id = A.compact(dim, index)
+    A_compact_dense = sparse_matrix_to_dense(A_compact)
+
+    A_dense = sparse_matrix_to_dense(A)
+
+    for i in range(shape[dim]):
+        if dim == 0:
+            row = list(A_dense[i, :].nonzero().reshape(-1))
+        else:
+            row = list(A_dense[:, i].nonzero().reshape(-1))
+        if (i not in list(ans_idx)) and len(row) > 0:
+            ans_idx.append(i)
+    if len(ans_idx):
+        ans_idx = torch.tensor(ans_idx).to(ctx)
+    A_dense_select = sparse_matrix_to_dense(A.index_select(dim, ans_idx))
+
+    assert A_compact_dense.shape == A_dense_select.shape
+    assert torch.allclose(A_compact_dense, A_dense_select)
+    assert torch.allclose(ans_idx, ret_id)