send_ue_recv

paddle.geometric. send_ue_recv ( x, y, src_index, dst_index, message_op='add', reduce_op='sum', out_size=None, name=None ) [source]

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.

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]]
Parameters
  • 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) – Different message ops for x and e, including add, sub, mul, div.

  • 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 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

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")
# Outputs: [[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)
# Outputs: [[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")
# Outputs: [[1., 3., 4.], [4., 10., 12.], [0., 0., 0.]]