Add _sparse_diag_cat().

src/icosagon/fastconv.py

@@ -5,6 +5,37 @@ from .data import Data
 from .trainprep import PreparedData
 import torch
 from .weights import init_glorot
+from .normalize import _sparse_coo_tensor
+
+
+def _sparse_diag_cat(matrices: List[torch.Tensor]):
+    if len(matrices) == 0:
+        raise ValueError('The list of matrices must be non-empty')
+
+    if not all(m.is_sparse for m in matrices):
+        raise ValueError('All matrices must be sparse')
+
+    if not all(len(m.shape) == 2 for m in matrices):
+        raise ValueError('All matrices must be 2D')
+
+    indices = []
+    values = []
+    row_offset = 0
+    col_offset = 0
+
+    for m in matrices:
+        ind = m._indices().clone()
+        ind[0] += row_offset
+        ind[1] += col_offset
+        indices.append(ind)
+        values.append(m._values())
+        row_offset += m.shape[0]
+        col_offset += m.shape[1]
+
+    indices = torch.cat(indices, dim=1)
+    values = torch.cat(values)
+
+    return _sparse_coo_tensor(indices, values, size=(row_offset, col_offset))
 
 
 def _cat(matrices: List[torch.Tensor]):
@@ -79,7 +110,7 @@ class FastConvLayer(torch.nn.Module):
         output_dim: List[int],
         data: Union[Data, PreparedData],
         keep_prob: float = 1.,
-        rel_activation: Callable[[torch.Tensor], torch.Tensor] = lambda x: x
+        rel_activation: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
         layer_activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.relu,
         **kwargs):
 

tests/icosagon/test_fastconv.py

@@ -0,0 +1,17 @@
+from icosagon.fastconv import _sparse_diag_cat
+import torch
+
+
+def test_sparse_diag_cat_01():
+    matrices = [ torch.rand(5, 10).round() for _ in range(7) ]
+    ground_truth = torch.zeros(35, 70)
+    ground_truth[0:5, 0:10] = matrices[0]
+    ground_truth[5:10, 10:20] = matrices[1]
+    ground_truth[10:15, 20:30] = matrices[2]
+    ground_truth[15:20, 30:40] = matrices[3]
+    ground_truth[20:25, 40:50] = matrices[4]
+    ground_truth[25:30, 50:60] = matrices[5]
+    ground_truth[30:35, 60:70] = matrices[6]
+    res = _sparse_diag_cat([ m.to_sparse() for m in matrices ])
+    res = res.to_dense()
+    assert torch.all(res == ground_truth)