Source code for dgl.udf

"""User-defined function related data structures."""
from __future__ import absolute_import


class EdgeBatch(object):
    """The class that can represent a batch of edges.

    Parameters
    ----------
    graph : DGLGraph
        Graph object.
    eid : Tensor
        Edge IDs.
    etype : (str, str, str)
        Edge type.
    src_data : dict[str, Tensor]
        Src node features.
    edge_data : dict[str, Tensor]
        Edge features.
    dst_data : dict[str, Tensor]
        Dst node features.
    """

    def __init__(self, graph, eid, etype, src_data, edge_data, dst_data):
        self._graph = graph
        self._eid = eid
        self._etype = etype
        self._src_data = src_data
        self._edge_data = edge_data
        self._dst_data = dst_data

    @property
    def src(self):
        """Return a view of the source node features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a node 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 retrieves the source node features for edges.
        >>> def edge_udf(edges):
        >>>     # edges.src['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'src': edges.src['h']}

        >>> # Copy features from source nodes to edges.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['src']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing, which is equivalent to
        >>> # dgl.function.copy_u.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('src', 'h'))
        >>> g.ndata['h']
        tensor([[1.],
                [2.]])
        """
        return self._src_data

    @property
    def dst(self):
        """Return a view of the destination node features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a node feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.tensor([[0.], [1.]])

        >>> # Define a UDF that retrieves the destination node features for
        >>> # edges.
        >>> def edge_udf(edges):
        >>>     # edges.dst['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'dst': edges.dst['h']}

        >>> # Copy features from destination nodes to edges.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['dst']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('dst', 'h'))
        >>> g.ndata['h']
        tensor([[0.],
                [2.]])
        """
        return self._dst_data

    @property
    def data(self):
        """Return a view of the edge features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set an edge feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.edata['h'] = torch.tensor([[1.], [1.], [1.]])

        >>> # Define a UDF that retrieves the feature 'h' for all edges.
        >>> def edge_udf(edges):
        >>>     # edges.data['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'data': edges.data['h']}

        >>> # Make a copy of the feature with name 'data'.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['data']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing, which is equivalent to
        >>> # dgl.function.copy_e.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('data', 'h'))
        >>> g.ndata['h']
        tensor([[1.],
                [2.]])
        """
        return self._edge_data

[docs] def edges(self): """Return the edges in the batch. Returns ------- (U, V, EID) : (Tensor, Tensor, Tensor) The edges in the batch. For each :math:`i`, :math:`(U[i], V[i])` is an edge from :math:`U[i]` to :math:`V[i]` with ID :math:`EID[i]`. Examples -------- The following example uses PyTorch backend. >>> import dgl >>> import torch >>> # Instantiate a graph. >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0]))) >>> # Define a UDF that retrieves and concatenates the end nodes of the >>> # edges. >>> def edge_udf(edges): >>> src, dst, _ = edges.edges() >>> return {'uv': torch.stack([src, dst], dim=1).float()} >>> # Create a feature 'uv' with the end nodes of the edges. >>> g.apply_edges(edge_udf) >>> g.edata['uv'] tensor([[0., 1.], [1., 1.], [1., 0.]]) >>> # Use edge UDF in message passing. >>> import dgl.function as fn >>> g.update_all(edge_udf, fn.sum('uv', 'h')) >>> g.ndata['h'] tensor([[1., 0.], [1., 2.]]) """ u, v = self._graph.find_edges(self._eid, etype=self.canonical_etype) return u, v, self._eid
[docs] def batch_size(self): """Return the number of edges in the batch. Returns ------- int Examples -------- The following example uses PyTorch backend. >>> import dgl >>> import torch >>> # Instantiate a graph. >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0]))) >>> # Define a UDF that returns one for each edge. >>> def edge_udf(edges): >>> return {'h': torch.ones(edges.batch_size(), 1)} >>> # Creates a feature 'h'. >>> g.apply_edges(edge_udf) >>> g.edata['h'] tensor([[1.], [1.], [1.]]) >>> # Use edge UDF in message passing. >>> import dgl.function as fn >>> g.update_all(edge_udf, fn.sum('h', 'h')) >>> g.ndata['h'] tensor([[1.], [2.]]) """ return len(self._eid)
def __len__(self): """Return the number of edges in this edge batch. Returns ------- int """ return self.batch_size() @property def canonical_etype(self): """Return the canonical edge type (i.e. triplet of source, edge, and destination node type) for this edge batch.""" return self._etype class NodeBatch(object): """The class to represent a batch of nodes. Parameters ---------- graph : DGLGraph Graph object. nodes : Tensor Node ids. ntype : str, optional The node type of this node batch, data : dict[str, Tensor] Node feature data. msgs : dict[str, Tensor], optional Messages data. """ def __init__(self, graph, nodes, ntype, data, msgs=None): self._graph = graph self._nodes = nodes self._ntype = ntype self._data = data self._msgs = msgs @property def data(self): """Return a view of the node features for the nodes in the batch. 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 feature for each node. >>> def node_udf(nodes): >>> # nodes.data['h'] is a tensor of shape (N, 1), >>> # 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.data['h'] + 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([[2.], [3.]]) """ return self._data @property def mailbox(self): """Return a view of the messages received. 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 feature for each node. >>> def node_udf(nodes): >>> # nodes.data['h'] is a tensor of shape (N, 1), >>> # 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.data['h'] + 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([[2.], [3.]]) """ return self._msgs
[docs] def nodes(self): """Return the nodes in the batch. Returns ------- NID : Tensor The IDs of the nodes in the batch. :math:`NID[i]` gives the ID of the i-th node. 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.]]) """ return self._nodes
[docs] def batch_size(self): """Return the number of nodes in the batch. Returns ------- int Examples -------- The following example uses PyTorch backend. >>> import dgl >>> import torch >>> # Instantiate a graph. >>> 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 for >>> # each node and increments the result by 1. >>> def node_udf(nodes): >>> return {'h': torch.ones(nodes.batch_size(), 1) >>> + 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([[2.], [3.]]) """ return len(self._nodes)
def __len__(self): """Return the number of nodes in this node batch. Returns ------- int """ return self.batch_size() @property def ntype(self): """Return the node type of this node batch, if available.""" return self._ntype