chore: import upstream snapshot with attribution

python/paddle/geometric/__init__.py

@@ -0,0 +1,32 @@
+# 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 .math import segment_max, segment_mean, segment_min, segment_sum
+from .message_passing import send_u_recv, send_ue_recv, send_uv
+from .reindex import reindex_graph, reindex_heter_graph
+from .sampling import sample_neighbors, weighted_sample_neighbors
+
+__all__ = [
+    'send_u_recv',
+    'send_ue_recv',
+    'send_uv',
+    'segment_sum',
+    'segment_mean',
+    'segment_min',
+    'segment_max',
+    'reindex_graph',
+    'reindex_heter_graph',
+    'sample_neighbors',
+    'weighted_sample_neighbors',
+]

python/paddle/geometric/math.py

@@ -0,0 +1,266 @@
+# 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
+
+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__ = []
+
+
+def segment_sum(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment Sum Operator.
+
+    This operator sums the elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that $out_i = \\sum_{j} data_{j}$
+    where sum is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (Tensor): A tensor, available data type float32, float64, int32, int64, float16.
+        segment_ids (Tensor): A 1-D tensor, which have the same size
+                            with the first dimension of input data.
+                            Available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - output (Tensor), the reduced result.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.geometric.segment_sum(data, segment_ids)
+            >>> print(out.numpy())
+            [[4. 4. 4.]
+             [4. 5. 6.]]
+
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "SUM")
+    else:
+        check_variable_and_dtype(
+            data,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16", "uint16"),
+            "segment_pool",
+        )
+        check_variable_and_dtype(
+            segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+        )
+
+        helper = LayerHelper("segment_sum", **locals())
+        out = helper.create_variable_for_type_inference(dtype=data.dtype)
+        summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+        helper.append_op(
+            type="segment_pool",
+            inputs={"X": data, "SegmentIds": segment_ids},
+            outputs={"Out": out, "SummedIds": summed_ids},
+            attrs={"pooltype": "SUM"},
+        )
+        return out
+
+
+def segment_mean(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment mean Operator.
+
+    This operator calculate the mean value of input `data` which
+    with the same index in `segment_ids`.
+    It computes a tensor such that $out_i = \\frac{1}{n_i}  \\sum_{j} data[j]$
+    where sum is over j such that 'segment_ids[j] == i' and $n_i$ is the number
+    of all index 'segment_ids[j] == i'.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64, float16.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - output (Tensor), the reduced result.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.geometric.segment_mean(data, segment_ids)
+            >>> print(out.numpy())
+            [[2. 2. 2.]
+             [4. 5. 6.]]
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "MEAN")
+    else:
+        check_variable_and_dtype(
+            data,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16", "uint16"),
+            "segment_pool",
+        )
+        check_variable_and_dtype(
+            segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+        )
+
+        helper = LayerHelper("segment_mean", **locals())
+        out = helper.create_variable_for_type_inference(dtype=data.dtype)
+        summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+        helper.append_op(
+            type="segment_pool",
+            inputs={"X": data, "SegmentIds": segment_ids},
+            outputs={"Out": out, "SummedIds": summed_ids},
+            attrs={"pooltype": "MEAN"},
+        )
+        return out
+
+
+def segment_min(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment min operator.
+
+    This operator calculate the minimum elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that $out_i = \\min_{j} data_{j}$
+    where min is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64, float16.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - output (Tensor), the reduced result.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.geometric.segment_min(data, segment_ids)
+            >>> print(out.numpy())
+            [[1. 2. 1.]
+             [4. 5. 6.]]
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "MIN")
+    else:
+        check_variable_and_dtype(
+            data,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16", "uint16"),
+            "segment_pool",
+        )
+        check_variable_and_dtype(
+            segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+        )
+
+        helper = LayerHelper("segment_min", **locals())
+        out = helper.create_variable_for_type_inference(dtype=data.dtype)
+        summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+        helper.append_op(
+            type="segment_pool",
+            inputs={"X": data, "SegmentIds": segment_ids},
+            outputs={"Out": out, "SummedIds": summed_ids},
+            attrs={"pooltype": "MIN"},
+        )
+        return out
+
+
+def segment_max(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment max operator.
+
+    This operator calculate the maximum elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that $out_i = \\max_{j} data_{j}$
+    where max is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64, float16.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - output (Tensor), the reduced result.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.geometric.segment_max(data, segment_ids)
+            >>> print(out.numpy())
+            [[3. 2. 3.]
+             [4. 5. 6.]]
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "MAX")
+    else:
+        check_variable_and_dtype(
+            data,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16", "uint16"),
+            "segment_pool",
+        )
+        check_variable_and_dtype(
+            segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+        )
+
+        helper = LayerHelper("segment_max", **locals())
+        out = helper.create_variable_for_type_inference(dtype=data.dtype)
+        summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+        helper.append_op(
+            type="segment_pool",
+            inputs={"X": data, "SegmentIds": segment_ids},
+            outputs={"Out": out, "SummedIds": summed_ids},
+            attrs={"pooltype": "MAX"},
+        )
+        return out

python/paddle/geometric/message_passing/__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 .send_recv import send_u_recv, send_ue_recv, send_uv  # noqa: F401
+
+__all__ = []

python/paddle/geometric/message_passing/send_recv.py

@@ -0,0 +1,538 @@
+#   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, TypeAlias
+
+import numpy as np
+
+from paddle import _C_ops
+from paddle.base.data_feeder import (
+    check_dtype,
+    check_type,
+    check_variable_and_dtype,
+)
+from paddle.base.framework import Variable
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+from .utils import (
+    convert_out_size_to_list,
+    get_out_size_tensor_inputs,
+    reshape_lhs_rhs,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+    _ReduceOp: TypeAlias = Literal[
+        "sum",
+        "mean",
+        "max",
+        "min",
+    ]
+    _MessageOp: TypeAlias = Literal[
+        "add",
+        "sub",
+        "mul",
+        "div",
+    ]
+__all__ = []
+
+
+def send_u_recv(
+    x: Tensor,
+    src_index: Tensor,
+    dst_index: Tensor,
+    reduce_op: _ReduceOp = "sum",
+    out_size: int | Tensor | None = None,
+    name: str | None = None,
+) -> Tensor:
+    """
+    Graph Learning message passing api.
+
+    This api is mainly used in Graph Learning domain, and the main purpose is to reduce intermediate memory
+    consumption in the process of message passing. Take `x` as the input tensor, we first use `src_index`
+    to gather the corresponding data, and then use `dst_index` to update the corresponding position of output tensor
+    in different reduce ops, like sum, mean, max, or min. Besides, we can use `out_size` to set necessary output shape.
+
+    .. code-block:: text
+
+           Given:
+
+           x = [[0, 2, 3],
+                [1, 4, 5],
+                [2, 6, 7]]
+
+           src_index = [0, 1, 2, 0]
+
+           dst_index = [1, 2, 1, 0]
+
+           reduce_op = "sum"
+
+           out_size = None
+
+           Then:
+
+           out = [[0, 2, 3],
+                  [2, 8, 10],
+                  [1, 4, 5]]
+
+    Args:
+        x (Tensor): The input tensor, and the available data type is float32, float64, int32, int64.
+                    And we support float16 in gpu version.
+        src_index (Tensor): An 1-D tensor, and the available data type is int32, int64.
+        dst_index (Tensor): An 1-D tensor, and should have the same shape as `src_index`.
+                            The available data type is int32, int64.
+        reduce_op (str): Different reduce ops, including `sum`, `mean`, `max`, `min`.
+                         Default value is `sum`.
+        out_size (int|Tensor|None): We can set `out_size` to get necessary output shape. If not set or
+                                    out_size is smaller or equal to 0, then this input will not be used.
+                                    Otherwise, `out_size` should be equal with or larger than
+                                    max(dst_index) + 1.
+        name (str, optional): Name for the operation (optional, default is None).
+                              For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - out (Tensor), the output tensor, should have the same shape and same dtype as input tensor `x`.
+          If `out_size` is set correctly, then it should have the same shape as `x` except the 0th dimension.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [1, 2], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out = paddle.geometric.send_u_recv(x, src_index, dst_index, reduce_op="sum")
+            >>> print(out.numpy())
+            [[ 0. 2. 3.]
+             [ 2. 8. 10.]
+             [ 1. 4. 5.]]
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out_size = paddle.max(dst_index) + 1
+            >>> out = paddle.geometric.send_u_recv(x, src_index, dst_index, reduce_op="sum", out_size=out_size)
+            >>> print(out.numpy())
+            [[ 0. 2. 3.]
+             [ 2. 8. 10.]]
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out = paddle.geometric.send_u_recv(x, src_index, dst_index, reduce_op="sum")
+            >>> print(out.numpy())
+            [[ 0. 2. 3.]
+             [ 2. 8. 10.]
+             [ 0. 0. 0.]]
+
+    """
+
+    if reduce_op not in ["sum", "mean", "max", "min"]:
+        raise ValueError(
+            f"reduce_op should be `sum`, `mean`, `max` or `min`, but received {reduce_op}"
+        )
+
+    # TODO(daisiming): Should we add judgement for out_size: max(dst_index) + 1.
+
+    if in_dynamic_or_pir_mode():
+        out_size = convert_out_size_to_list(out_size, 'graph_send_recv')
+        return _C_ops.send_u_recv(
+            x, src_index, dst_index, reduce_op.upper(), out_size
+        )
+    else:
+        check_variable_and_dtype(
+            x,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16"),
+            "graph_send_recv",
+        )
+        check_variable_and_dtype(
+            src_index, "Src_index", ("int32", "int64"), "graph_send_recv"
+        )
+        check_variable_and_dtype(
+            dst_index, "Dst_index", ("int32", "int64"), "graph_send_recv"
+        )
+        if out_size:
+            check_type(
+                out_size,
+                'out_size',
+                (int, np.int32, np.int64, Variable),
+                'graph_send_recv',
+            )
+        if isinstance(out_size, Variable):
+            check_dtype(
+                out_size.dtype,
+                'out_size',
+                ['int32', 'int64'],
+                'graph_send_recv',
+            )
+
+        helper = LayerHelper("send_u_recv", **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+        dst_count = helper.create_variable_for_type_inference(
+            dtype="int32", stop_gradient=True
+        )
+
+        inputs = {"X": x, "Src_index": src_index, "Dst_index": dst_index}
+        attrs = {"reduce_op": reduce_op.upper()}
+        get_out_size_tensor_inputs(
+            inputs=inputs,
+            attrs=attrs,
+            out_size=out_size,
+            op_type='graph_send_recv',
+        )
+
+        helper.append_op(
+            type="graph_send_recv",
+            inputs=inputs,
+            outputs={"Out": out, "Dst_count": dst_count},
+            attrs=attrs,
+        )
+        return out
+
+
+def send_ue_recv(
+    x: Tensor,
+    y: Tensor,
+    src_index: Tensor,
+    dst_index: Tensor,
+    message_op: _MessageOp = "add",
+    reduce_op: _ReduceOp = "sum",
+    out_size: int | Tensor | None = None,
+    name: str | None = None,
+) -> Tensor:
+    """
+
+    Graph Learning message passing api.
+
+    This api is mainly used in Graph Learning domain, and the main purpose is to reduce intermediate memory
+    consumption in the process of message passing. Take `x` as the input tensor, we first use `src_index`
+    to gather the corresponding data, after computing with `y` in different message ops like add/sub/mul/div, then use `dst_index` to
+    update the corresponding position of output tensor in different reduce ops, like sum, mean, max, or min.
+    Besides, we can use `out_size` to set necessary output shape.
+
+    .. code-block:: text
+
+           Given:
+
+           x = [[0, 2, 3],
+                [1, 4, 5],
+                [2, 6, 7]]
+
+           y = [1, 1, 1]
+
+           src_index = [0, 1, 2, 0]
+
+           dst_index = [1, 2, 1, 0]
+
+           message_op = "add"
+
+           reduce_op = "sum"
+
+           out_size = None
+
+           Then:
+
+           out = [[1, 3, 4],
+                  [4, 10, 12],
+                  [2, 5, 6]]
+
+    Args:
+        x (Tensor): The input node feature tensor, and the available data type is float32, float64, int32, int64.
+                    And we support float16 in gpu version.
+        y (Tensor): The input edge feature tensor, and the available data type is float32, float64, int32, int64.
+                    And we support float16 in gpu version.
+        src_index (Tensor): An 1-D tensor, and the available data type is int32, int64.
+        dst_index (Tensor): An 1-D tensor, and should have the same shape as `src_index`.
+                            The available data type is int32, int64.
+        message_op (str, optional): Different message ops for x and e, including `add`, `sub`, `mul`, `div`.
+        reduce_op (str, optional): Different reduce ops, including `sum`, `mean`, `max`, `min`.
+                         Default value is `sum`.
+        out_size (int|Tensor, optional): We can set `out_size` to get necessary output shape. If not set or
+                                    out_size is smaller or equal to 0, then this input will not be used.
+                                    Otherwise, `out_size` should be equal with or larger than
+                                    max(dst_index) + 1. Default value 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 (Tensor), the output tensor, should have the same shape and same dtype as input tensor `x`.
+          If `out_size` is set correctly, then it should have the same shape as `x` except the 0th dimension.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> y = paddle.to_tensor([1, 1, 1, 1], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [1, 2], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out = paddle.geometric.send_ue_recv(x, y, src_index, dst_index, message_op="add", reduce_op="sum")
+            >>> print(out.numpy())
+            [[ 1. 3. 4.]
+             [ 4. 10. 12.]
+             [ 2. 5. 6.]]
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> y = paddle.to_tensor([1, 1, 1], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out_size = paddle.max(dst_index) + 1
+            >>> out = paddle.geometric.send_ue_recv(x, y, src_index, dst_index, message_op="add", reduce_op="sum", out_size=out_size)
+            >>> print(out.numpy())
+            [[ 1. 3. 4.]
+             [ 4. 10. 12.]]
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> y = paddle.to_tensor([1, 1, 1], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index, dst_index = indexes[:, 0], indexes[:, 1]
+            >>> out = paddle.geometric.send_ue_recv(x, y, src_index, dst_index, message_op="add", reduce_op="sum")
+            >>> print(out.numpy())
+            [[ 1. 3. 4.]
+             [ 4. 10. 12.]
+             [ 0. 0. 0.]]
+
+    """
+
+    if message_op not in ["add", "sub", "mul", "div"]:
+        raise ValueError(
+            f"message_op should be `add`, `sub`, `mul`, `div`, but received {message_op}"
+        )
+
+    if reduce_op not in ["sum", "mean", "max", "min"]:
+        raise ValueError(
+            f"reduce_op should be `sum`, `mean`, `max` or `min`, but received {reduce_op}"
+        )
+
+    x, y = reshape_lhs_rhs(x, y)
+
+    if message_op == 'sub':
+        message_op = 'add'
+        y = -y
+    if message_op == "div":
+        message_op = 'mul'
+        y = 1.0 / (y + 1e-12)
+
+    # TODO(daisiming): Should we add judgement for out_size: max(dst_index) + 1.
+
+    if in_dynamic_or_pir_mode():
+        out_size = convert_out_size_to_list(out_size, 'graph_send_ue_recv')
+        return _C_ops.send_ue_recv(
+            x,
+            y,
+            src_index,
+            dst_index,
+            message_op.upper(),
+            reduce_op.upper(),
+            out_size,
+        )
+    else:
+        check_variable_and_dtype(
+            x,
+            "X",
+            ("float32", "float64", "int32", "int64", "float16"),
+            "graph_send_ue_recv",
+        )
+        check_variable_and_dtype(
+            y,
+            "Y",
+            ("float32", "float64", "int32", "int64", "float16"),
+            "graph_send_ue_recv",
+        )
+        check_variable_and_dtype(
+            src_index, "Src_index", ("int32", "int64"), "graph_send_ue_recv"
+        )
+        check_variable_and_dtype(
+            dst_index, "Dst_index", ("int32", "int64"), "graph_send_ue_recv"
+        )
+        if out_size:
+            check_type(
+                out_size,
+                'out_size',
+                (int, np.int32, np.int64, Variable),
+                'graph_send_ue_recv',
+            )
+        if isinstance(out_size, Variable):
+            check_dtype(
+                out_size.dtype,
+                'out_size',
+                ['int32', 'int64'],
+                'graph_send_ue_recv',
+            )
+
+        helper = LayerHelper("send_ue_recv", **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+        dst_count = helper.create_variable_for_type_inference(
+            dtype="int32", stop_gradient=True
+        )
+
+        inputs = {
+            "X": x,
+            "Y": y,
+            "Src_index": src_index,
+            "Dst_index": dst_index,
+        }
+        attrs = {
+            "message_op": message_op.upper(),
+            "reduce_op": reduce_op.upper(),
+        }
+        get_out_size_tensor_inputs(
+            inputs=inputs,
+            attrs=attrs,
+            out_size=out_size,
+            op_type='graph_send_ue_recv',
+        )
+
+        helper.append_op(
+            type="graph_send_ue_recv",
+            inputs=inputs,
+            outputs={"Out": out, "Dst_count": dst_count},
+            attrs=attrs,
+        )
+        return out
+
+
+def send_uv(
+    x: Tensor,
+    y: Tensor,
+    src_index: Tensor,
+    dst_index: Tensor,
+    message_op: _MessageOp = "add",
+    name: str | None = None,
+) -> Tensor:
+    """
+
+    Graph Learning message passing api.
+
+    This api is mainly used in Graph Learning domain, and the main purpose is to reduce intermediate memory
+    consumption in the process of message passing. Take `x` as the source node feature tensor, take `y` as
+    the destination node feature tensor. Then we use `src_index` and `dst_index` to gather the corresponding data,
+    and then compute the edge features in different message_ops like `add`, `sub`, `mul`, `div`.
+
+    .. code-block:: text
+
+           Given:
+
+           x = [[0, 2, 3],
+                [1, 4, 5],
+                [2, 6, 7]]
+
+           y = [[0, 1, 2],
+                [2, 3, 4],
+                [4, 5, 6]]
+
+           src_index = [0, 1, 2, 0]
+
+           dst_index = [1, 2, 1, 0]
+
+           message_op = "add"
+
+           Then:
+
+           out = [[2, 5, 7],
+                  [5, 9, 11],
+                  [4, 9, 11],
+                  [0, 3, 5]]
+
+    Args:
+        x (Tensor): The source node feature tensor, and the available data type is float32, float64, int32, int64. And we support float16 in gpu version.
+        y (Tensor): The destination node feature tensor, and the available data type is float32, float64, int32, int64. And we support float16 in gpu version.
+        src_index (Tensor): An 1-D tensor, and the available data type is int32, int64.
+        dst_index (Tensor): An 1-D tensor, and should have the same shape as `src_index`.
+                            The available data type is int32, int64.
+        message_op (str): Different message ops for x and y, including `add`, `sub`, `mul` and `div`.
+        name (str, optional): Name for the operation (optional, default is None).
+                              For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        - out (Tensor), the output tensor.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = paddle.to_tensor([[0, 2, 3], [1, 4, 5], [2, 6, 7]], dtype="float32")
+            >>> y = paddle.to_tensor([[0, 1, 2], [2, 3, 4], [4, 5, 6]], dtype="float32")
+            >>> indexes = paddle.to_tensor([[0, 1], [1, 2], [2, 1], [0, 0]], dtype="int32")
+            >>> src_index = indexes[:, 0]
+            >>> dst_index = indexes[:, 1]
+            >>> out = paddle.geometric.send_uv(x, y, src_index, dst_index, message_op="add")
+            >>> print(out.numpy())
+            [[ 2. 5. 7.]
+             [ 5. 9. 11.]
+             [ 4. 9. 11.]
+             [ 0. 3. 5.]]
+
+    """
+
+    if message_op not in ['add', 'sub', 'mul', 'div']:
+        raise ValueError(
+            f"message_op should be `add`, `sub`, `mul`, `div`, but received {message_op}"
+        )
+
+    x, y = reshape_lhs_rhs(x, y)
+
+    if message_op == 'sub':
+        message_op = 'add'
+        y = -y
+    if message_op == 'div':
+        message_op = 'mul'
+        y = 1.0 / (y + 1e-12)
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.send_uv(x, y, src_index, dst_index, message_op.upper())
+    else:
+        helper = LayerHelper("graph_send_uv", **locals())
+        check_variable_and_dtype(
+            x,
+            'x',
+            ['int32', 'int64', 'float32', 'float64', 'float16'],
+            'graph_send_uv',
+        )
+        check_variable_and_dtype(
+            y,
+            'y',
+            ['int32', 'int64', 'float32', 'float64', 'float16'],
+            'graph_send_uv',
+        )
+        check_variable_and_dtype(
+            src_index, 'src_index', ['int32', 'int64'], 'graph_send_uv'
+        )
+        check_variable_and_dtype(
+            dst_index, 'dst_index', ['int32', 'int64'], 'graph_send_uv'
+        )
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+        inputs = {
+            'x': x,
+            'y': y,
+            'src_index': src_index,
+            'dst_index': dst_index,
+        }
+        attrs = {'message_op': message_op.upper()}
+        helper.append_op(
+            type="graph_send_uv",
+            inputs=inputs,
+            attrs=attrs,
+            outputs={"out": out},
+        )
+        return out

python/paddle/geometric/message_passing/utils.py

@@ -0,0 +1,111 @@
+#   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.
+
+import numpy as np
+
+import paddle
+from paddle.base.data_feeder import check_dtype, convert_dtype
+from paddle.base.framework import Variable
+
+
+def convert_out_size_to_list(out_size, op_type):
+    """
+    Convert out_size(int, np.int32, np.int64, Variable) to list
+    in imperative mode.
+    """
+    if out_size is None:
+        out_size = [0]
+    elif isinstance(out_size, (int, np.int32, np.int64)):
+        out_size = [out_size]
+    elif isinstance(out_size, (Variable, paddle.pir.Value)):
+        out_size.stop_gradient = True
+        check_dtype(
+            out_size.dtype,
+            'out_size',
+            ['int32', 'int64'],
+            'op_type',
+            '(When type of out_size in' + op_type + ' is Variable.)',
+        )
+        if convert_dtype(out_size.dtype) == 'int64':
+            out_size = paddle.cast(out_size, 'int32')
+    else:
+        out_size = [int(out_size)]
+    return out_size
+
+
+def get_out_size_tensor_inputs(inputs, attrs, out_size, op_type):
+    """
+    Convert out_size(int, np.int32, np.int64, Variable) to inputs
+    and attrs in static graph mode.
+    """
+    if out_size is None:
+        attrs['out_size'] = [0]
+    elif isinstance(out_size, (int, np.int32, np.int64)):
+        attrs['out_size'] = [out_size]
+    elif isinstance(out_size, Variable):
+        out_size.stop_gradient = True
+        check_dtype(
+            out_size.dtype,
+            'out_size',
+            ['int32', 'int64'],
+            'op_type',
+            '(When type of out_size in' + op_type + ' is Variable.)',
+        )
+        if convert_dtype(out_size.dtype) == 'int64':
+            out_size = paddle.cast(out_size, 'int32')
+        inputs["Out_size"] = out_size
+    else:
+        raise TypeError("Out_size only supports Variable or int.")
+
+
+def reshape_lhs_rhs(x, y):
+    """
+    Expand dims to ensure there will be no broadcasting issues with different
+    number of dimensions.
+    """
+    if len(x.shape) == 1:
+        x = paddle.reshape(x, [-1, 1])
+    if len(y.shape) == 1:
+        y = paddle.reshape(y, [-1, 1])
+
+    x_shape = paddle.shape(x)
+    y_shape = paddle.shape(y)
+    if len(x.shape) != len(y.shape):
+        max_ndims = max(len(x.shape), len(y.shape))
+        x_pad_ndims = max_ndims - len(x.shape)
+        y_pad_ndims = max_ndims - len(y.shape)
+        new_x_shape = (
+            [
+                x_shape[0],
+            ]
+            + [
+                1,
+            ]
+            * x_pad_ndims
+            + list(x_shape[1:])
+        )
+        new_y_shape = (
+            [
+                y_shape[0],
+            ]
+            + [
+                1,
+            ]
+            * y_pad_ndims
+            + list(y_shape[1:])
+        )
+        x = paddle.reshape(x, new_x_shape)
+        y = paddle.reshape(y, new_y_shape)
+
+    return x, y

python/paddle/geometric/reindex.py

@@ -0,0 +1,289 @@
+#   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
+
+import paddle
+from paddle import _C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.framework import Variable
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+
+__all__ = []
+
+
+def reindex_graph(
+    x: Tensor,
+    neighbors: Tensor,
+    count: Tensor,
+    value_buffer: Tensor | None = None,
+    index_buffer: Tensor | None = None,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """
+
+    Reindex Graph API.
+
+    This API is mainly used in Graph Learning domain, which should be used
+    in conjunction with `paddle.geometric.sample_neighbors` API. And the main purpose
+    is to reindex the ids information of the input nodes, and return the
+    corresponding graph edges after reindex.
+
+    Take input nodes x = [0, 1, 2] as an example. If we have neighbors = [8, 9, 0, 4, 7, 6, 7], and count = [2, 3, 2],
+    then we know that the neighbors of 0 is [8, 9], the neighbors of 1 is [0, 4, 7], and the neighbors of 2 is [6, 7].
+    Then after graph_reindex, we will have 3 different outputs: reindex_src: [3, 4, 0, 5, 6, 7, 6], reindex_dst: [0, 0, 1, 1, 1, 2, 2]
+    and out_nodes: [0, 1, 2, 8, 9, 4, 7, 6]. We can see that the numbers in `reindex_src` and `reindex_dst` is the corresponding index
+    of nodes in `out_nodes`.
+
+    Note:
+        The number in x should be unique, otherwise it would cause potential errors. We will reindex all the nodes from 0.
+
+    Args:
+        x (Tensor): The input nodes which we sample neighbors for. The available
+                    data type is int32, int64.
+        neighbors (Tensor): The neighbors of the input nodes `x`. The data type
+                            should be the same with `x`.
+        count (Tensor): The neighbor count of the input nodes `x`. And the
+                        data type should be int32.
+        value_buffer (Tensor, optional): Value buffer for hashtable. The data type should be int32,
+                                    and should be filled with -1. Only useful for gpu version. Default is None.
+        index_buffer (Tensor, optional): Index buffer for hashtable. The data type should be int32,
+                                    and should be filled with -1. Only useful for gpu version.
+                                    `value_buffer` and `index_buffer` should be both not None
+                                    if you want to speed up by using hashtable buffer. 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:
+        - reindex_src (Tensor), the source node index of graph edges after reindex.
+
+        - reindex_dst (Tensor), the destination node index of graph edges after reindex.
+
+        - out_nodes (Tensor), the index of unique input nodes and neighbors before reindex, where we put the input nodes `x` in the front, and put neighbor nodes in the back.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = [0, 1, 2]
+            >>> neighbors = [8, 9, 0, 4, 7, 6, 7]
+            >>> count = [2, 3, 2]
+            >>> x = paddle.to_tensor(x, dtype="int64")
+            >>> neighbors = paddle.to_tensor(neighbors, dtype="int64")
+            >>> count = paddle.to_tensor(count, dtype="int32")
+            >>> reindex_src, reindex_dst, out_nodes = paddle.geometric.reindex_graph(x, neighbors, count)
+            >>> print(reindex_src.numpy())
+            [3 4 0 5 6 7 6]
+            >>> print(reindex_dst.numpy())
+            [0 0 1 1 1 2 2]
+            >>> print(out_nodes.numpy())
+            [0 1 2 8 9 4 7 6]
+
+    """
+    use_buffer_hashtable = (
+        True if value_buffer is not None and index_buffer is not None else False
+    )
+
+    if in_dynamic_or_pir_mode():
+        reindex_src, reindex_dst, out_nodes = _C_ops.reindex_graph(
+            x,
+            neighbors,
+            count,
+            value_buffer,
+            index_buffer,
+        )
+        return reindex_src, reindex_dst, out_nodes
+
+    check_variable_and_dtype(x, "X", ("int32", "int64"), "graph_reindex")
+    check_variable_and_dtype(
+        neighbors, "Neighbors", ("int32", "int64"), "graph_reindex"
+    )
+    check_variable_and_dtype(count, "Count", ("int32"), "graph_reindex")
+
+    if use_buffer_hashtable:
+        check_variable_and_dtype(
+            value_buffer, "HashTable_Value", ("int32"), "graph_reindex"
+        )
+        check_variable_and_dtype(
+            index_buffer, "HashTable_Index", ("int32"), "graph_reindex"
+        )
+
+    helper = LayerHelper("reindex_graph", **locals())
+    reindex_src = helper.create_variable_for_type_inference(dtype=x.dtype)
+    reindex_dst = helper.create_variable_for_type_inference(dtype=x.dtype)
+    out_nodes = helper.create_variable_for_type_inference(dtype=x.dtype)
+    helper.append_op(
+        type="graph_reindex",
+        inputs={
+            "X": x,
+            "Neighbors": neighbors,
+            "Count": count,
+            "HashTable_Value": value_buffer if use_buffer_hashtable else None,
+            "HashTable_Index": index_buffer if use_buffer_hashtable else None,
+        },
+        outputs={
+            "Reindex_Src": reindex_src,
+            "Reindex_Dst": reindex_dst,
+            "Out_Nodes": out_nodes,
+        },
+    )
+    return reindex_src, reindex_dst, out_nodes
+
+
+def reindex_heter_graph(
+    x: Tensor,
+    neighbors: Sequence[Tensor],
+    count: Sequence[Tensor],
+    value_buffer: Tensor | None = None,
+    index_buffer: Tensor | None = None,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """
+
+    Reindex HeterGraph API.
+
+    This API is mainly used in Graph Learning domain, which should be used
+    in conjunction with `paddle.geometric.sample_neighbors` API. And the main purpose
+    is to reindex the ids information of the input nodes, and return the
+    corresponding graph edges after reindex.
+
+    Take input nodes x = [0, 1, 2] as an example. For graph A, suppose we have neighbors = [8, 9, 0, 4, 7, 6, 7], and count = [2, 3, 2],
+    then we know that the neighbors of 0 is [8, 9], the neighbors of 1 is [0, 4, 7], and the neighbors of 2 is [6, 7]. For graph B,
+    suppose we have neighbors = [0, 2, 3, 5, 1], and count = [1, 3, 1], then we know that the neighbors of 0 is [0], the neighbors of 1 is [2, 3, 5],
+    and the neighbors of 3 is [1]. We will get following outputs: reindex_src: [3, 4, 0, 5, 6, 7, 6, 0, 2, 8, 9, 1], reindex_dst: [0, 0, 1, 1, 1, 2, 2, 0, 1, 1, 1, 2]
+    and out_nodes: [0, 1, 2, 8, 9, 4, 7, 6, 3, 5].
+
+    Note:
+        The number in x should be unique, otherwise it would cause potential errors. We support multi-edge-types neighbors reindexing in reindex_heter_graph api. We will reindex all the nodes from 0.
+
+    Args:
+        x (Tensor): The input nodes which we sample neighbors for. The available
+                    data type is int32, int64.
+        neighbors (list|tuple): The neighbors of the input nodes `x` from different graphs.
+                                The data type should be the same with `x`.
+        count (list|tuple): The neighbor counts of the input nodes `x` from different graphs.
+                            And the data type should be int32.
+        value_buffer (Tensor, optional): Value buffer for hashtable. The data type should be int32,
+                                    and should be filled with -1. Only useful for gpu version. Default is None.
+        index_buffer (Tensor, optional): Index buffer for hashtable. The data type should be int32,
+                                    and should be filled with -1. Only useful for gpu version.
+                                    `value_buffer` and `index_buffer` should be both not None
+                                    if you want to speed up by using hashtable buffer. 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:
+        - reindex_src (Tensor), the source node index of graph edges after reindex.
+
+        - reindex_dst (Tensor), the destination node index of graph edges after reindex.
+
+        - out_nodes (Tensor), the index of unique input nodes and neighbors before reindex,
+                              where we put the input nodes `x` in the front, and put neighbor
+                              nodes in the back.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+
+            >>> x = [0, 1, 2]
+            >>> neighbors_a = [8, 9, 0, 4, 7, 6, 7]
+            >>> count_a = [2, 3, 2]
+            >>> x = paddle.to_tensor(x, dtype="int64")
+            >>> neighbors_a = paddle.to_tensor(neighbors_a, dtype="int64")
+            >>> count_a = paddle.to_tensor(count_a, dtype="int32")
+            >>> neighbors_b = [0, 2, 3, 5, 1]
+            >>> count_b = [1, 3, 1]
+            >>> neighbors_b = paddle.to_tensor(neighbors_b, dtype="int64")
+            >>> count_b = paddle.to_tensor(count_b, dtype="int32")
+            >>> neighbors = [neighbors_a, neighbors_b]
+            >>> count = [count_a, count_b]
+            >>> reindex_src, reindex_dst, out_nodes = paddle.geometric.reindex_heter_graph(x, neighbors, count)
+            >>> print(reindex_src.numpy())
+            [3 4 0 5 6 7 6 0 2 8 9 1]
+            >>> print(reindex_dst.numpy())
+            [0 0 1 1 1 2 2 0 1 1 1 2]
+            >>> print(out_nodes.numpy())
+            [0 1 2 8 9 4 7 6 3 5]
+
+    """
+    use_buffer_hashtable = (
+        True if value_buffer is not None and index_buffer is not None else False
+    )
+
+    if in_dynamic_or_pir_mode():
+        neighbors = paddle.concat(neighbors, axis=0)
+        count = paddle.concat(count, axis=0)
+        reindex_src, reindex_dst, out_nodes = _C_ops.reindex_graph(
+            x,
+            neighbors,
+            count,
+            value_buffer,
+            index_buffer,
+        )
+        return reindex_src, reindex_dst, out_nodes
+
+    if isinstance(neighbors, Variable):
+        neighbors = [neighbors]
+    if isinstance(count, Variable):
+        count = [count]
+
+    neighbors = paddle.concat(neighbors, axis=0)
+    count = paddle.concat(count, axis=0)
+
+    check_variable_and_dtype(x, "X", ("int32", "int64"), "heter_graph_reindex")
+    check_variable_and_dtype(
+        neighbors, "Neighbors", ("int32", "int64"), "graph_reindex"
+    )
+    check_variable_and_dtype(count, "Count", ("int32"), "graph_reindex")
+
+    if use_buffer_hashtable:
+        check_variable_and_dtype(
+            value_buffer, "HashTable_Value", ("int32"), "graph_reindex"
+        )
+        check_variable_and_dtype(
+            index_buffer, "HashTable_Index", ("int32"), "graph_reindex"
+        )
+
+    helper = LayerHelper("reindex_heter_graph", **locals())
+    reindex_src = helper.create_variable_for_type_inference(dtype=x.dtype)
+    reindex_dst = helper.create_variable_for_type_inference(dtype=x.dtype)
+    out_nodes = helper.create_variable_for_type_inference(dtype=x.dtype)
+    neighbors = paddle.concat(neighbors, axis=0)
+    count = paddle.concat(count, axis=0)
+    helper.append_op(
+        type="graph_reindex",
+        inputs={
+            "X": x,
+            "Neighbors": neighbors,
+            "Count": count,
+            "HashTable_Value": value_buffer if use_buffer_hashtable else None,
+            "HashTable_Index": index_buffer if use_buffer_hashtable else None,
+        },
+        outputs={
+            "Reindex_Src": reindex_src,
+            "Reindex_Dst": reindex_dst,
+            "Out_Nodes": out_nodes,
+        },
+    )
+    return reindex_src, reindex_dst, out_nodes

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