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.
>>> import dgl >>> import torch
>>> # Instantiate a graph and set a feature 'h'. >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0]))) >>> g.ndata['h'] = torch.ones(2, 1)
>>> # 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)}
>>> # Use node UDF in message passing. >>> import dgl.function as fn >>> g.update_all(fn.copy_u('h', 'm'), node_udf) >>> g.ndata['h'] tensor([[1.], [3.]])