chore: import upstream snapshot with attribution

python/paddle/geometric/sampling/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 .neighbors import sample_neighbors, weighted_sample_neighbors  # noqa: F401
+
+__all__ = []

python/paddle/geometric/sampling/neighbors.py

@@ -0,0 +1,405 @@
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
+from paddle import _C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+__all__ = []
+
+
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[True] = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor, Tensor]: ...
+
+
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[False] = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: bool = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: ...
+
+
+def sample_neighbors(
+    row,
+    colptr,
+    input_nodes,
+    sample_size=-1,
+    eids=None,
+    return_eids=False,
+    perm_buffer=None,
+    name=None,
+):
+    """
+
+    Graph Sample Neighbors API.
+
+    This API is mainly used in Graph Learning domain, and the main purpose is to
+    provide high performance of graph sampling method. For example, we get the
+    CSC(Compressed Sparse Column) format of the input graph edges as `row` and
+    `colptr`, so as to convert graph data into a suitable format for sampling.
+    `input_nodes` means the nodes we need to sample neighbors, and `sample_sizes`
+    means the number of neighbors and number of layers we want to sample.
+
+    Besides, we support fisher-yates sampling in GPU version.
+
+    Args:
+        row (Tensor): One of the components of the CSC format of the input graph, and
+                      the shape should be [num_edges, 1] or [num_edges]. The available
+                      data type is int32, int64.
+        colptr (Tensor): One of the components of the CSC format of the input graph,
+                         and the shape should be [num_nodes + 1, 1] or [num_nodes + 1].
+                         The data type should be the same with `row`.
+        input_nodes (Tensor): The input nodes we need to sample neighbors for, and the
+                              data type should be the same with `row`.
+        sample_size (int, optional): The number of neighbors we need to sample. Default value is -1,
+                           which means returning all the neighbors of the input nodes.
+        eids (Tensor, optional): The eid information of the input graph. If return_eids is True,
+                            then `eids` should not be None. The data type should be the
+                            same with `row`. Default is None.
+        return_eids (bool, optional): Whether to return eid information of sample edges. Default is False.
+        perm_buffer (Tensor, optional): Permutation buffer for fisher-yates sampling. If `use_perm_buffer`
+                              is True, then `perm_buffer` should not be None. The data type should
+                              be the same with `row`. If not None, we will use fiser-yates sampling
+                              to speed up. Only useful for gpu version. Default is None.
+        name (str, optional): Name for the operation (optional, default is None).
+                              For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - out_neighbors (Tensor), the sample neighbors of the input nodes.
+
+        - out_count (Tensor), the number of sampling neighbors of each input node, and the shape
+          should be the same with `input_nodes`.
+
+        - out_eids (Tensor), if `return_eids` is True, we will return the eid information of the
+          sample edges.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> # edges: (3, 0), (7, 0), (0, 1), (9, 1), (1, 2), (4, 3), (2, 4),
+            >>> #        (9, 5), (3, 5), (9, 6), (1, 6), (9, 8), (7, 8)
+            >>> row = [3, 7, 0, 9, 1, 4, 2, 9, 3, 9, 1, 9, 7]
+            >>> colptr = [0, 2, 4, 5, 6, 7, 9, 11, 11, 13, 13]
+            >>> nodes = [0, 8, 1, 2]
+            >>> sample_size = 2
+            >>> row = paddle.to_tensor(row, dtype="int64")
+            >>> colptr = paddle.to_tensor(colptr, dtype="int64")
+            >>> nodes = paddle.to_tensor(nodes, dtype="int64")
+            >>> out_neighbors, out_count = paddle.geometric.sample_neighbors(
+            ...     row, colptr, nodes, sample_size=sample_size, return_eids=False
+            ... )
+
+    """
+
+    if return_eids:
+        if eids is None:
+            raise ValueError(
+                "`eids` should not be None if `return_eids` is True."
+            )
+
+    use_perm_buffer = True if perm_buffer is not None else False
+
+    if in_dynamic_or_pir_mode():
+        (
+            out_neighbors,
+            out_count,
+            out_eids,
+        ) = _C_ops.graph_sample_neighbors(
+            row,
+            colptr,
+            input_nodes,
+            eids,
+            perm_buffer,
+            sample_size,
+            return_eids,
+            use_perm_buffer,
+        )
+        if return_eids:
+            return out_neighbors, out_count, out_eids
+        return out_neighbors, out_count
+
+    check_variable_and_dtype(
+        row, "Row", ("int32", "int64"), "graph_sample_neighbors"
+    )
+    check_variable_and_dtype(
+        colptr, "Col_Ptr", ("int32", "int64"), "graph_sample_neighbors"
+    )
+    check_variable_and_dtype(
+        input_nodes, "X", ("int32", "int64"), "graph_sample_neighbors"
+    )
+    if return_eids:
+        check_variable_and_dtype(
+            eids, "Eids", ("int32", "int64"), "graph_sample_neighbors"
+        )
+    if use_perm_buffer:
+        check_variable_and_dtype(
+            perm_buffer,
+            "Perm_Buffer",
+            ("int32", "int64"),
+            "graph_sample_neighbors",
+        )
+
+    helper = LayerHelper("sample_neighbors", **locals())
+    out_neighbors = helper.create_variable_for_type_inference(dtype=row.dtype)
+    out_count = helper.create_variable_for_type_inference(dtype=row.dtype)
+    out_eids = helper.create_variable_for_type_inference(dtype=row.dtype)
+    helper.append_op(
+        type="graph_sample_neighbors",
+        inputs={
+            "Row": row,
+            "Col_Ptr": colptr,
+            "X": input_nodes,
+            "Eids": eids if return_eids else None,
+            "Perm_Buffer": perm_buffer if use_perm_buffer else None,
+        },
+        outputs={
+            "Out": out_neighbors,
+            "Out_Count": out_count,
+            "Out_Eids": out_eids,
+        },
+        attrs={
+            "sample_size": sample_size,
+            "return_eids": return_eids,
+            "flag_perm_buffer": use_perm_buffer,
+        },
+    )
+    if return_eids:
+        return out_neighbors, out_count, out_eids
+    return out_neighbors, out_count
+
+
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[True] = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor, Tensor]: ...
+
+
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[False] = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor]: ...
+
+
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: bool = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: ...
+
+
+def weighted_sample_neighbors(
+    row,
+    colptr,
+    edge_weight,
+    input_nodes,
+    sample_size=-1,
+    eids=None,
+    return_eids=False,
+    name=None,
+):
+    """
+    Graph Weighted Sample Neighbors API.
+
+    This API is mainly used in Graph Learning domain, and the main purpose is to
+    provide high performance of graph weighted-sampling method. For example, we get the
+    CSC(Compressed Sparse Column) format of the input graph edges as `row` and
+    `colptr`, so as to convert graph data into a suitable format for sampling, and the
+    input `edge_weight` should also match the CSC format. Besides, `input_nodes` means
+    the nodes we need to sample neighbors, and `sample_sizes` means the number of neighbors
+    and number of layers we want to sample. This API will finally return the weighted sampled
+    neighbors, and the probability of being selected as a neighbor is related to its weight,
+    with higher weight and higher probability.
+
+    Args:
+        row (Tensor): One of the components of the CSC format of the input graph, and
+                      the shape should be [num_edges, 1] or [num_edges]. The available
+                      data type is int32, int64.
+        colptr (Tensor): One of the components of the CSC format of the input graph,
+                         and the shape should be [num_nodes + 1, 1] or [num_nodes + 1].
+                         The data type should be the same with `row`.
+        edge_weight (Tensor): The edge weight of the CSC format graph edges. And the shape
+                              should be [num_edges, 1] or [num_edges]. The available data
+                              type is float32.
+        input_nodes (Tensor): The input nodes we need to sample neighbors for, and the
+                              data type should be the same with `row`.
+        sample_size (int, optional): The number of neighbors we need to sample. Default value is -1,
+                           which means returning all the neighbors of the input nodes.
+        eids (Tensor, optional): The eid information of the input graph. If return_eids is True,
+                            then `eids` should not be None. The data type should be the
+                            same with `row`. Default is None.
+        return_eids (bool, optional): Whether to return eid information of sample edges. Default is False.
+        name (str, optional): Name for the operation (optional, default is None).
+                              For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - out_neighbors (Tensor), the sample neighbors of the input nodes.
+
+        - out_count (Tensor), the number of sampling neighbors of each input node, and the shape
+          should be the same with `input_nodes`.
+
+        - out_eids (Tensor), if `return_eids` is True, we will return the eid information of the
+          sample edges.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> # edges: (3, 0), (7, 0), (0, 1), (9, 1), (1, 2), (4, 3), (2, 4),
+            >>> #        (9, 5), (3, 5), (9, 6), (1, 6), (9, 8), (7, 8)
+            >>> row = [3, 7, 0, 9, 1, 4, 2, 9, 3, 9, 1, 9, 7]
+            >>> colptr = [0, 2, 4, 5, 6, 7, 9, 11, 11, 13, 13]
+            >>> weight = [0.1, 0.5, 0.2, 0.5, 0.9, 1.9, 2.0, 2.1, 0.01, 0.9, 0, 12, 0.59, 0.67]
+            >>> nodes = [0, 8, 1, 2]
+            >>> sample_size = 2
+            >>> row = paddle.to_tensor(row, dtype="int64")
+            >>> colptr = paddle.to_tensor(colptr, dtype="int64")
+            >>> weight = paddle.to_tensor(weight, dtype="float32")
+            >>> nodes = paddle.to_tensor(nodes, dtype="int64")
+            >>> out_neighbors, out_count = paddle.geometric.weighted_sample_neighbors(
+            ...     row, colptr, weight, nodes, sample_size=sample_size, return_eids=False
+            ... )
+
+    """
+
+    if return_eids:
+        if eids is None:
+            raise ValueError(
+                "`eids` should not be None if `return_eids` is True."
+            )
+
+    if in_dynamic_or_pir_mode():
+        (
+            out_neighbors,
+            out_count,
+            out_eids,
+        ) = _C_ops.weighted_sample_neighbors(
+            row,
+            colptr,
+            edge_weight,
+            input_nodes,
+            eids,
+            sample_size,
+            return_eids,
+        )
+        if return_eids:
+            return out_neighbors, out_count, out_eids
+        return out_neighbors, out_count
+
+    check_variable_and_dtype(
+        row, "row", ("int32", "int64"), "weighted_sample_neighbors"
+    )
+    check_variable_and_dtype(
+        colptr, "colptr", ("int32", "int64"), "weighted_sample_neighbors"
+    )
+    check_variable_and_dtype(
+        edge_weight,
+        "edge_weight",
+        ("float32"),
+        "weighted_sample_neighbors",
+    )
+    check_variable_and_dtype(
+        input_nodes,
+        "input_nodes",
+        ("int32", "int64"),
+        "weighted_sample_neighbors",
+    )
+    if return_eids:
+        check_variable_and_dtype(
+            eids, "eids", ("int32", "int64"), "weighted_sample_neighbors"
+        )
+
+    helper = LayerHelper("weighted_sample_neighbors", **locals())
+    out_neighbors = helper.create_variable_for_type_inference(dtype=row.dtype)
+    out_count = helper.create_variable_for_type_inference(dtype=row.dtype)
+    out_eids = helper.create_variable_for_type_inference(dtype=row.dtype)
+    helper.append_op(
+        type="weighted_sample_neighbors",
+        inputs={
+            "row": row,
+            "colptr": colptr,
+            "edge_weight": edge_weight,
+            "input_nodes": input_nodes,
+            "eids": eids if return_eids else None,
+        },
+        outputs={
+            "out_neighbors": out_neighbors,
+            "out_count": out_count,
+            "out_eids": out_eids,
+        },
+        attrs={
+            "sample_size": sample_size,
+            "return_eids": return_eids,
+        },
+    )
+    if return_eids:
+        return out_neighbors, out_count, out_eids
+    return out_neighbors, out_count