support empty adj

tensorkit/backend/pytorch_/sparse.py

@@ -25,7 +25,7 @@
 
     # sparse tensor operations
     'coalesce', 'is_coalesced', 'get_indices', 'get_values',
-    'rank', 'length', 'shape', 'get_dtype', 'get_device',
+    'rank', 'length', 'shape', 'value_count', 'get_dtype', 'get_device',
     'to_dtype', 'to_device', 'eye', 'reduce_sum', 'matmul',
 
     # sparse tensor grad utilities
@@ -208,6 +208,13 @@ def shape(input: Tensor) -> List[int]:
     return list(input.shape)
 
 
+@sparse_jit
+def value_count(input: Tensor) -> List[int]:
+    if not input.is_coalesced():
+        input = input.coalesce()
+    return input.indices().shape[1]
+
+
 @sparse_jit
 def get_dtype(input: Tensor) -> str:
     if input.dtype == torch.float32:

tensorkit/gnn/adj/tensor_ops.py

@@ -54,6 +54,16 @@ def _col_div(m: T.SparseTensor, v: T.Tensor) -> T.SparseTensor:
     return T.sparse.make_sparse(indices, v, T.sparse.shape(m))
 
 
+def _compute_degree(adj_matrix: T.SparseTensor) -> T.Tensor:
+    if T.sparse.value_count(adj_matrix) == 0:
+        size = T.sparse.length(adj_matrix)
+        degree = T.zeros([size], dtype=T.sparse.get_dtype(adj_matrix),
+                         device=T.sparse.get_device(adj_matrix))
+    else:
+        degree = T.sparse.to_dense(T.sparse.reduce_sum(adj_matrix, axis=-1))
+    return degree
+
+
 @T.jit_ignore
 def normalize_adj(adj_matrix: T.SparseTensor,
                   undirected: bool = False,
@@ -63,7 +73,7 @@ def normalize_adj(adj_matrix: T.SparseTensor,
 
     # normalize the adj matrix
     # TODO: do we need to check whether or not all degrees are non-negative?
-    degree = T.sparse.to_dense(T.sparse.reduce_sum(adj_matrix, axis=-1))
+    degree = _compute_degree(adj_matrix)
     if epsilon != 0.:
         degree = T.clip_left(degree, epsilon)
 
@@ -94,7 +104,7 @@ def normalize_partitioned_adj(adj_matrices: List[T.SparseTensor],
 
     # normalize the adj matrices
     degree = T.add_n([
-        T.sparse.to_dense(T.sparse.reduce_sum(adj, axis=-1))
+        _compute_degree(adj)
         for adj in adj_matrices
     ])
     if epsilon != 0.:

tests/gnn/adj/test_tensor_ops.py

@@ -46,7 +46,15 @@ def D(t):
         def G(d, y):
             return np.dot(D(d), y)
 
-        x_list = [make_random_adj_matrix(node_count) for _ in range(3)]
+        empty_adj = T.sparse.from_dense(T.zeros([node_count, node_count]))
+        self.assertEqual(
+            T.shape(T.sparse.get_indices(empty_adj, coord_first=True))[1],
+            0
+        )
+        x_list = (
+            [make_random_adj_matrix(node_count) for _ in range(3)] +
+            [empty_adj]
+        )
         y_list = [T.sparse.to_numpy(x) for x in x_list]
         d_list = [np.maximum(np.sum(y, axis=-1), eps) for y in y_list]
         d_sum = sum(d_list, 0.)

tests/tensor/test_sparse.py

@@ -59,6 +59,7 @@ def h(x, y, t_):
                     **kwargs
                 )
                 g(x, y)
+                self.assertEqual(T.sparse.value_count(y), len(row))
 
             # from_dense
             y = T.sparse.from_dense(
@@ -230,6 +231,19 @@ def test_reduce_sum(self):
                     rtol=1e-4, atol=1e-6,
                 )
 
+    def test_no_element_sparse_tensor(self):
+        for coord_first in [True, False]:
+            # construct the no-element sparse tensor
+            indices_shape = [2, 0] if coord_first else [0, 2]
+            x = T.sparse.make_sparse(
+                T.zeros(indices_shape, dtype=T.index_dtype),
+                T.zeros([0], dtype=T.float32),
+                coord_first=coord_first,
+                shape=[3, 4],
+            )
+            assert_allclose(x, np.zeros([3, 4]))
+            self.assertEqual(T.sparse.value_count(x), 0)
+
     def test_matmul(self):
         indices = T.as_tensor(np.random.randint(0, 50, size=[2, 200]))
         values = T.random.randn([200])