chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
@@ -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__ = []
@@ -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
@@ -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