[cugraph-pyg] Add TransformerConv and support for bipartite node features in GATConv

ci/test_python.sh

@@ -179,6 +179,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
  --channel "${PYTHON_CHANNEL}" \
  libcugraph \
  pylibcugraph \
+ pylibcugraphops \
  cugraph \
  cugraph-pyg
 

python/cugraph-pyg/cugraph_pyg/nn/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .conv import *

python/cugraph-pyg/cugraph_pyg/nn/conv/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .gat_conv import GATConv
+from .gatv2_conv import GATv2Conv
+from .transformer_conv import TransformerConv
+
+__all__ = [
+ "GATConv",
+ "GATv2Conv",
+ "TransformerConv",
+]

python/cugraph-pyg/cugraph_pyg/nn/conv/base.py

@@ -0,0 +1,196 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Any, Optional, Tuple, Union
+
+from cugraph.utilities.utils import import_optional
+
+torch = import_optional("torch")
+torch_geometric = import_optional("torch_geometric")
+
+try: # pragma: no cover
+ from pylibcugraphops.pytorch import (
+ BipartiteCSC,
+ SampledCSC,
+ SampledHeteroCSC,
+ StaticCSC,
+ StaticHeteroCSC,
+ )
+
+ HAS_PYLIBCUGRAPHOPS = True
+except ImportError:
+ HAS_PYLIBCUGRAPHOPS = False
+
+
+class BaseConv(torch.nn.Module): # pragma: no cover
+ r"""An abstract base class for implementing cugraph-ops message passing layers."""
+
+ def __init__(self):
+ super().__init__()
+
+ if HAS_PYLIBCUGRAPHOPS is False:
+ raise ModuleNotFoundError(
+ f"'{self.__class__.__name__}' requires " f"'pylibcugraphops>=23.04'"
+ )
+
+ def reset_parameters(self):
+ r"""Resets all learnable parameters of the module."""
+ pass
+
+ @staticmethod
+ def to_csc(
+ edge_index: torch.Tensor,
+ size: Optional[Tuple[int, int]] = None,
+ edge_attr: Optional[torch.Tensor] = None,
+ ) -> Union[
+ Tuple[torch.Tensor, torch.Tensor, int],
+ Tuple[Tuple[torch.Tensor, torch.Tensor, int], torch.Tensor],
+ ]:
+ r"""Returns a CSC representation of an :obj:`edge_index` tensor to be
+ used as input to cugraph-ops conv layers.
+
+ Args:
+ edge_index (torch.Tensor): The edge indices.
+ size ((int, int), optional). The shape of :obj:`edge_index` in each
+ dimension. (default: :obj:`None`)
+ edge_attr (torch.Tensor, optional): The edge features.
+ (default: :obj:`None`)
+ """
+ if size is None:
+ warnings.warn(
+ f"Inferring the graph size from 'edge_index' causes "
+ f"a decline in performance and does not work for "
+ f"bipartite graphs. To suppress this warning, pass "
+ f"the 'size' explicitly in '{__name__}.to_csc()'."
+ )
+ num_src_nodes = num_dst_nodes = int(edge_index.max()) + 1
+ else:
+ num_src_nodes, num_dst_nodes = size
+
+ row, col = edge_index
+ col, perm = torch_geometric.utils.index_sort(col, max_value=num_dst_nodes)
+ row = row[perm]
+
+ colptr = torch_geometric.utils.sparse.index2ptr(col, num_dst_nodes)
+
+ if edge_attr is not None:
+ return (row, colptr, num_src_nodes), edge_attr[perm]
+
+ return row, colptr, num_src_nodes
+
+ def get_cugraph(
+ self,
+ csc: Tuple[torch.Tensor, torch.Tensor, int],
+ bipartite: bool = False,
+ max_num_neighbors: Optional[int] = None,
+ ) -> Any:
+ r"""Constructs a :obj:`cugraph-ops` graph object from CSC representation.
+ Supports both bipartite and non-bipartite graphs.
+
+ Args:
+ csc ((torch.Tensor, torch.Tensor, int)): A tuple containing the CSC
+ representation of a graph, given as a tuple of
+ :obj:`(row, colptr, num_src_nodes)`. Use the
+ :meth:`to_csc` method to convert an :obj:`edge_index`
+ representation to the desired format.
+ bipartite (bool): If set to :obj:`True`, will create the bipartite
+ structure in cugraph-ops. (default: :obj:`False`)
+ max_num_neighbors (int, optional): The maximum number of neighbors
+ of a target node. It is only effective when operating in a
+ bipartite graph. When not given, will be computed on-the-fly,
+ leading to slightly worse performance. (default: :obj:`None`)
+ """
+ row, colptr, num_src_nodes = csc
+
+ if not row.is_cuda:
+ raise RuntimeError(
+ f"'{self.__class__.__name__}' requires GPU-"
+ f"based processing (got CPU tensor)"
+ )
+
+ if bipartite:
+ return BipartiteCSC(colptr, row, num_src_nodes)
+
+ if num_src_nodes != colptr.numel() - 1:
+ if max_num_neighbors is None:
+ max_num_neighbors = int((colptr[1:] - colptr[:-1]).max())
+
+ return SampledCSC(colptr, row, max_num_neighbors, num_src_nodes)
+
+ return StaticCSC(colptr, row)
+
+ def get_typed_cugraph(
+ self,
+ csc: Tuple[torch.Tensor, torch.Tensor, int],
+ edge_type: torch.Tensor,
+ num_edge_types: Optional[int] = None,
+ bipartite: bool = False,
+ max_num_neighbors: Optional[int] = None,
+ ) -> Any:
+ r"""Constructs a typed :obj:`cugraph` graph object from a CSC
+ representation where each edge corresponds to a given edge type.
+ Supports both bipartite and non-bipartite graphs.
+
+ Args:
+ csc ((torch.Tensor, torch.Tensor, int)): A tuple containing the CSC
+ representation of a graph, given as a tuple of
+ :obj:`(row, colptr, num_src_nodes)`. Use the
+ :meth:`to_csc` method to convert an :obj:`edge_index`
+ representation to the desired format.
+ edge_type (torch.Tensor): The edge type.
+ num_edge_types (int, optional): The maximum number of edge types.
+ When not given, will be computed on-the-fly, leading to
+ slightly worse performance. (default: :obj:`None`)
+ bipartite (bool): If set to :obj:`True`, will create the bipartite
+ structure in cugraph-ops. (default: :obj:`False`)
+ max_num_neighbors (int, optional): The maximum number of neighbors
+ of a target node. It is only effective when operating in a
+ bipartite graph. When not given, will be computed on-the-fly,
+ leading to slightly worse performance. (default: :obj:`None`)
+ """
+ if num_edge_types is None:
+ num_edge_types = int(edge_type.max()) + 1
+
+ row, colptr, num_src_nodes = csc
+ edge_type = edge_type.int()
+
+ if bipartite:
+ raise NotImplementedError
+
+ if num_src_nodes != colptr.numel() - 1:
+ if max_num_neighbors is None:
+ max_num_neighbors = int((colptr[1:] - colptr[:-1]).max())
+
+ return SampledHeteroCSC(
+ colptr, row, edge_type, max_num_neighbors, num_src_nodes, num_edge_types
+ )
+
+ return StaticHeteroCSC(colptr, row, edge_type, num_edge_types)
+
+ def forward(
+ self,
+ x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+ csc: Tuple[torch.Tensor, torch.Tensor, int],
+ ) -> torch.Tensor:
+ r"""Runs the forward pass of the module.
+
+ Args:
+ x (torch.Tensor): The node features.
+ csc ((torch.Tensor, torch.Tensor, int)): A tuple containing the CSC
+ representation of a graph, given as a tuple of
+ :obj:`(row, colptr, num_src_nodes)`. Use the
+ :meth:`to_csc` method to convert an :obj:`edge_index`
+ representation to the desired format.
+ """
+ raise NotImplementedError