Add spmm and is_torch_sparse_tensor

test/utils/test_torch_sparse_tensor.py

@@ -0,0 +1,12 @@
+import torch
+from torch_sparse import SparseTensor
+
+from torch_geometric.utils import is_torch_sparse_tensor
+
+
+def test_is_torch_sparse_tensor():
+    x = torch.randn(5, 5)
+
+    assert not is_torch_sparse_tensor(x)
+    assert not is_torch_sparse_tensor(SparseTensor.from_dense(x))
+    assert is_torch_sparse_tensor(x.to_sparse())

torch_geometric/nn/functional/__init__.py

@@ -1,6 +1,7 @@
 from .bro import bro
 from .gini import gini
+from .spmm import spmm
 
-__all__ = ['bro', 'gini']
+__all__ = ['bro', 'gini', 'spmm']
 
 classes = __all__

torch_geometric/nn/functional/spmm.py

@@ -0,0 +1,22 @@
+import torch
+from torch import Tensor
+from torch_sparse import SparseTensor
+
+from torch_geometric.utils import is_torch_sparse_tensor
+
+
+def spmm(src: SparseTensor, other: Tensor, reduce: str = "sum") -> Tensor:
+    assert reduce in {"sum", "mean", "max", "min"}
+    if isinstance(src, SparseTensor):
+        return src.spmm(other, reduce)
+    if not is_torch_sparse_tensor(src):
+        raise ValueError("`src` must be a torch_sparse SparseTensor "
+                         f"or a PyTorch SparseTensor, but got {type(src)}.")
+    if reduce == "sum":
+        return torch.sparse.mm(src, other)
+    elif reduce == "mean":
+        # TODO: Support `mean` reduction for PyTorch SparseTensor
+        raise NotImplementedError
+    else:
+        raise NotImplementedError("`max` and `min` reduction are not supported"
+                                  "for PyTorch SparseTensor input.")

torch_geometric/utils/__init__.py

@@ -36,6 +36,7 @@
                                 structured_negative_sampling_feasible)
 from .train_test_split_edges import train_test_split_edges
 from .scatter import scatter
+from .torch_sparse_tensor import is_torch_sparse_tensor
 
 __all__ = [
     'degree',
@@ -95,6 +96,7 @@
     'structured_negative_sampling_feasible',
     'train_test_split_edges',
     'scatter',
+    'is_torch_sparse_tensor',
 ]
 
 classes = __all__

torch_geometric/utils/torch_sparse_tensor.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+
+def is_torch_sparse_tensor(x: Any) -> bool:
+    """Returns :obj:`True` if the input :obj:`x` is
+    PyTorch SparseTensor (COO or CSR format).
+
+    Args:
+        x (Any): The input object to be checked.
+    """
+    return getattr(x, 'is_sparse', False)