modify spmm to support multi-dimensional tensor

a simple modification to the code of spmm() can make it support also 3D tensor (might think as batched):
suppose dense matrix has dimensions (B, N, F), the nonzero has dimensions (B, num_edges), all graphs share the same topology (i.e., the same adjacency matrix)

torch_sparse/spmm.py

@@ -15,13 +15,13 @@ def spmm(index, value, m, n, matrix):
     :rtype: :class:`Tensor`
     """
 
-    assert n == matrix.size(0)
+    assert n == matrix.shape[-2]
 
     row, col = index
     matrix = matrix if matrix.dim() > 1 else matrix.unsqueeze(-1)
 
-    out = matrix[col]
+    out = matrix[..., col, :]
     out = out * value.unsqueeze(-1)
-    out = scatter_add(out, row, dim=0, dim_size=m)
+    out = scatter_add(out, row, dim=-2, dim_size=m)
 
     return out