def send(src, dst, nfeat, efeat, message_func): """Send message from src to dst. """ src_feat = op.RowReader(nfeat, src) dst_feat = op.RowReader(nfeat, dst) msg = message_func(src_feat, dst_feat, efeat) return msg
def recv(self, reduce_func, msg, recv_mode="dst"): """Recv message and aggregate the message by reduce_func The UDF reduce_func function should has the following format. .. code-block:: python def reduce_func(msg): ''' Args: msg: A LodTensor or a dictionary of LodTensor whose batch_size is equals to the number of unique dst nodes. Return: It should return a tensor with shape (batch_size, out_dims). The batch size should be the same as msg. ''' pass Args: msg: A tensor or a dictionary of tensor created by send function.. reduce_func: A callable UDF reduce function. Return: A tensor with shape (num_nodes, out_dims). The output for nodes with no message will be zeros. """ if not self._is_tensor: raise ValueError("You must call Graph.tensor()") if not isinstance(msg, dict): raise TypeError( "The input of msg should be a dict, but receives a %s" % (type(msg))) if not callable(reduce_func): raise TypeError("reduce_func should be callable") src, dst, eid = self.sorted_edges(sort_by=recv_mode) msg = op.RowReader(msg, eid) if recv_mode == "dst": uniq_ind, segment_ids = paddle.unique(dst, return_inverse=True) elif recv_mode == "src": uniq_ind, segment_ids = paddle.unique(src, return_inverse=True) bucketed_msg = Message(msg, segment_ids) output = reduce_func(bucketed_msg) output_dim = output.shape[-1] init_output = paddle.zeros(shape=[self._num_nodes, output_dim], dtype=output.dtype) final_output = scatter(init_output, uniq_ind, output) return final_output
def send( self, message_func, src_feat=None, dst_feat=None, edge_feat=None, node_feat=None, ): """Send message from all src nodes to dst nodes. The UDF message function should has the following format. .. code-block:: python def message_func(src_feat, dst_feat, edge_feat): ''' Args: src_feat: the node feat dict attached to the src nodes. dst_feat: the node feat dict attached to the dst nodes. edge_feat: the edge feat dict attached to the corresponding (src, dst) edges. Return: It should return a tensor or a dictionary of tensor. And each tensor should have a shape of (num_edges, dims). ''' return {'msg': src_feat['h']} Args: message_func: UDF function. src_feat: a dict {name: tensor,} to build src node feat dst_feat: a dict {name: tensor,} to build dst node feat node_feat: a dict {name: tensor,} to build both src and dst node feat edge_feat: a dict {name: tensor,} to build edge feat Return: A dictionary of tensor representing the message. Each of the values in the dictionary has a shape (num_edges, dim) which should be collected by :code:`recv` function. """ if self._is_tensor: if (src_feat is not None or dst_feat is not None) and node_feat is not None: raise ValueError( "Can not use src/dst feat and node feat at the same time") src_feat_temp = {} dst_feat_temp = {} if node_feat is not None: assert isinstance(node_feat, dict), "The input node_feat must be a dict" src_feat_temp.update(node_feat) dst_feat_temp.update(node_feat) else: if src_feat is not None: assert isinstance( src_feat, dict), "The input src_feat must be a dict" src_feat_temp.update(src_feat) if dst_feat is not None: assert isinstance( dst_feat, dict), "The input dst_feat must be a dict" dst_feat_temp.update(dst_feat) edge_feat_temp = {} if edge_feat is not None: assert isinstance(edge_feat, dict), "The input edge_feat must be a dict" edge_feat_temp.update(edge_feat) src = self.edges[:, 0] dst = self.edges[:, 1] src_feat = op.RowReader(src_feat_temp, src) dst_feat = op.RowReader(dst_feat_temp, dst) msg = message_func(src_feat, dst_feat, edge_feat_temp) if not isinstance(msg, dict): raise TypeError( "The outputs of the %s function is expected to be a dict, but got %s" \ % (message_func.__name__, type(msg))) return msg else: raise ValueError("You must call Graph.tensor() first")