graph_send_recv¶
- paddle.incubate. graph_send_recv ( x, src_index, dst_index, pool_type='sum', out_size=None, name=None ) [source]
-
Graph Learning Send_Recv combine operator.
This operator 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 pooling types, like sum, mean, max, or min. Besides, we can set out_size to get necessary output shape.
Given: X = [[0, 2, 3], [1, 4, 5], [2, 6, 7]] src_index = [0, 1, 2, 0] dst_index = [1, 2, 1, 0] pool_type = "sum" out_size = None Then: Out = [[0, 2, 3], [2, 8, 10], [1, 4, 5]]
- Parameters
-
x (Tensor) – The input tensor, and the available data type is float32, float64, int32, int64.
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.
pool_type (str) – The pooling types of graph_send_recv, 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 Name.
- Returns
-
- 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.
- Return type
-
out (Tensor)
Examples
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 = indexes[:, 0] dst_index = indexes[:, 1] out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum") # Outputs: [[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 = indexes[:, 0] dst_index = indexes[:, 1] out_size = paddle.max(dst_index) + 1 out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum", out_size=out_size) # Outputs: [[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 = indexes[:, 0] dst_index = indexes[:, 1] out = paddle.incubate.graph_send_recv(x, src_index, dst_index, pool_type="sum") # Outputs: [[0., 2., 3.], [2., 8., 10.], [0., 0., 0.]]