chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,415 @@
|
||||
"""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
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user