dgl.udf.NodeBatch.nodes¶
-
NodeBatch.
nodes
()[source]¶ Return the nodes in the batch.
- Returns
NID – The IDs of the nodes in the batch. NID[i] gives the ID of the i-th node.
- Return type
Tensor
Examples
The following example uses PyTorch backend.
>>> # Define a UDF that computes the sum of the messages received and the original ID >>> # for each node >>> def node_udf(nodes): >>> # nodes.nodes() is a tensor of shape (N), >>> # nodes.mailbox['m'] is a tensor of shape (N, D, 1), >>> # where N is the number of nodes in the batch, >>> # D is the number of messages received per node for this node batch >>> return {'h': nodes.nodes().unsqueeze(-1).float() + nodes.mailbox['m'].sum(1)}