chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
"""dgl operator module."""
|
||||
from .edge_softmax import *
|
||||
from .gather_mm import *
|
||||
from .sddmm import *
|
||||
from .segment import *
|
||||
from .spmm import *
|
||||
@@ -0,0 +1,153 @@
|
||||
"""dgl edge_softmax operator module."""
|
||||
from ..backend import (
|
||||
astype,
|
||||
edge_softmax as edge_softmax_internal,
|
||||
edge_softmax_hetero as edge_softmax_hetero_internal,
|
||||
)
|
||||
from ..base import ALL, is_all
|
||||
|
||||
__all__ = ["edge_softmax"]
|
||||
|
||||
|
||||
def edge_softmax(graph, logits, eids=ALL, norm_by="dst"):
|
||||
r"""Compute softmax over weights of incoming edges for every node.
|
||||
|
||||
For a node :math:`i`, edge softmax is an operation that computes
|
||||
|
||||
.. math::
|
||||
a_{ij} = \frac{\exp(z_{ij})}{\sum_{j\in\mathcal{N}(i)}\exp(z_{ij})}
|
||||
|
||||
where :math:`z_{ij}` is a signal of edge :math:`j\rightarrow i`, also
|
||||
called logits in the context of softmax. :math:`\mathcal{N}(i)` is
|
||||
the set of nodes that have an edge to :math:`i`.
|
||||
|
||||
By default edge softmax is normalized by destination nodes(i.e. :math:`ij`
|
||||
are incoming edges of `i` in the formula above). We also support edge
|
||||
softmax normalized by source nodes(i.e. :math:`ij` are outgoing edges of
|
||||
`i` in the formula). The former case corresponds to softmax in GAT and
|
||||
Transformer, and the latter case corresponds to softmax in Capsule network.
|
||||
An example of using edge softmax is in
|
||||
`Graph Attention Network <https://arxiv.org/pdf/1710.10903.pdf>`__ where
|
||||
the attention weights are computed with this operation.
|
||||
Other non-GNN examples using this are
|
||||
`Transformer <https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf>`__,
|
||||
`Capsule <https://arxiv.org/pdf/1710.09829.pdf>`__, etc.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph over which edge softmax will be performed.
|
||||
logits : torch.Tensor or dict of torch.Tensor
|
||||
The input edge feature. Heterogeneous graphs can have dict of tensors where
|
||||
each tensor stores the edge features of the corresponding relation type.
|
||||
eids : torch.Tensor or ALL, optional
|
||||
The IDs of the edges to apply edge softmax. If ALL, it will apply edge
|
||||
softmax to all edges in the graph. Default: ALL.
|
||||
norm_by : str, could be `src` or `dst`
|
||||
Normalized by source nodes or destination nodes. Default: `dst`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor or tuple of tensors
|
||||
Softmax value.
|
||||
|
||||
Notes
|
||||
-----
|
||||
* Input shape: :math:`(E, *, 1)` where * means any number of
|
||||
additional dimensions, :math:`E` equals the length of eids.
|
||||
If the `eids` is ALL, :math:`E` equals the number of edges in
|
||||
the graph.
|
||||
* Return shape: :math:`(E, *, 1)`
|
||||
|
||||
Examples on a homogeneous graph
|
||||
-------------------------------
|
||||
The following example uses PyTorch backend.
|
||||
|
||||
>>> from dgl.nn.functional import edge_softmax
|
||||
>>> import dgl
|
||||
>>> import torch as th
|
||||
|
||||
Create a :code:`DGLGraph` object and initialize its edge features.
|
||||
|
||||
>>> g = dgl.graph((th.tensor([0, 0, 0, 1, 1, 2]), th.tensor([0, 1, 2, 1, 2, 2])))
|
||||
>>> edata = th.ones(6, 1).float()
|
||||
>>> edata
|
||||
tensor([[1.],
|
||||
[1.],
|
||||
[1.],
|
||||
[1.],
|
||||
[1.],
|
||||
[1.]])
|
||||
|
||||
Apply edge softmax over g:
|
||||
|
||||
>>> edge_softmax(g, edata)
|
||||
tensor([[1.0000],
|
||||
[0.5000],
|
||||
[0.3333],
|
||||
[0.5000],
|
||||
[0.3333],
|
||||
[0.3333]])
|
||||
|
||||
Apply edge softmax over g normalized by source nodes:
|
||||
|
||||
>>> edge_softmax(g, edata, norm_by='src')
|
||||
tensor([[0.3333],
|
||||
[0.3333],
|
||||
[0.3333],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[1.0000]])
|
||||
|
||||
Apply edge softmax to first 4 edges of g:
|
||||
|
||||
>>> edge_softmax(g, edata[:4], th.Tensor([0,1,2,3]))
|
||||
tensor([[1.0000],
|
||||
[0.5000],
|
||||
[1.0000],
|
||||
[0.5000]])
|
||||
|
||||
|
||||
Examples on a heterogeneous graph
|
||||
---------------------------------
|
||||
|
||||
Create a heterogeneous graph and initialize its edge features.
|
||||
|
||||
>>> hg = dgl.heterograph({
|
||||
... ('user', 'follows', 'user'): ([0, 0, 1], [0, 1, 2]),
|
||||
... ('developer', 'develops', 'game'): ([0, 1], [0, 1])
|
||||
... })
|
||||
>>> edata_follows = th.ones(3, 1).float()
|
||||
>>> edata_develops = th.ones(2, 1).float()
|
||||
>>> edata_dict = {('user', 'follows', 'user'): edata_follows,
|
||||
... ('developer','develops', 'game'): edata_develops}
|
||||
|
||||
Apply edge softmax over hg normalized by source nodes:
|
||||
|
||||
>>> edge_softmax(hg, edata_dict, norm_by='src')
|
||||
{('developer', 'develops', 'game'): tensor([[1.],
|
||||
[1.]]), ('user', 'follows', 'user'): tensor([[0.5000],
|
||||
[0.5000],
|
||||
[1.0000]])}
|
||||
"""
|
||||
if not is_all(eids):
|
||||
eids = astype(eids, graph.idtype)
|
||||
if graph._graph.number_of_etypes() == 1:
|
||||
return edge_softmax_internal(
|
||||
graph._graph, logits, eids=eids, norm_by=norm_by
|
||||
)
|
||||
else:
|
||||
logits_list = [None] * graph._graph.number_of_etypes()
|
||||
logits = {graph.to_canonical_etype(k): v for k, v in logits.items()}
|
||||
for rel in graph.canonical_etypes:
|
||||
etid = graph.get_etype_id(rel)
|
||||
logits_list[etid] = logits[rel]
|
||||
logits_tuple = tuple(logits_list)
|
||||
score_tuple = edge_softmax_hetero_internal(
|
||||
graph._graph, eids, norm_by, *logits_tuple
|
||||
)
|
||||
score = {}
|
||||
for rel in graph.canonical_etypes:
|
||||
etid = graph.get_etype_id(rel)
|
||||
score[rel] = score_tuple[etid]
|
||||
return score
|
||||
@@ -0,0 +1,50 @@
|
||||
"""dgl gather_mm operator module."""
|
||||
from .. import backend as F
|
||||
|
||||
__all__ = ["gather_mm"]
|
||||
|
||||
|
||||
def gather_mm(a, b, *, idx_b):
|
||||
r"""Gather data according to the given indices and perform matrix multiplication.
|
||||
|
||||
Let the result tensor be ``c``, the operator conducts the following computation:
|
||||
|
||||
c[i] = a[i] @ b[idx_b[i]]
|
||||
, where len(c) == len(idx_b)
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a : Tensor
|
||||
A 2-D tensor of shape ``(N, D1)``
|
||||
b : Tensor
|
||||
A 3-D tensor of shape ``(R, D1, D2)``
|
||||
idx_b : Tensor, optional
|
||||
An 1-D integer tensor of shape ``(N,)``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor
|
||||
The output dense matrix of shape ``(N, D2)``
|
||||
"""
|
||||
N, D1 = F.shape(a)
|
||||
R, _, D2 = F.shape(b)
|
||||
if N > 1000000 or D1 > 8 or D2 > 8:
|
||||
# Use segment_mm for large workload
|
||||
import torch
|
||||
|
||||
sorted_idx_b, perm = torch.sort(idx_b)
|
||||
_, rev_perm = torch.sort(perm)
|
||||
sorted_a = torch.index_select(a, 0, perm)
|
||||
pos_l = torch.searchsorted(
|
||||
sorted_idx_b, torch.arange(R, device=a.device)
|
||||
)
|
||||
pos_r = torch.cat(
|
||||
[pos_l[1:], torch.tensor([len(idx_b)], device=a.device)]
|
||||
)
|
||||
seglen = (pos_r - pos_l).cpu() # XXX(minjie): cause device synchronize
|
||||
return torch.index_select(
|
||||
F.segment_mm(sorted_a, b, seglen), 0, rev_perm
|
||||
)
|
||||
else:
|
||||
return F.gather_mm(a, b, None, idx_b)
|
||||
@@ -0,0 +1,207 @@
|
||||
"""dgl sddmm operator module."""
|
||||
import sys
|
||||
from itertools import product
|
||||
|
||||
from .. import backend as F
|
||||
from ..backend import (
|
||||
gsddmm as gsddmm_internal,
|
||||
gsddmm_hetero as gsddmm_internal_hetero,
|
||||
)
|
||||
|
||||
__all__ = ["gsddmm", "copy_u", "copy_v", "copy_e"]
|
||||
|
||||
|
||||
def reshape_lhs_rhs(lhs_data, rhs_data):
|
||||
r"""Expand dims so that there will be no broadcasting issues with different
|
||||
number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
|
||||
that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
|
||||
(E, 5, 3, 4)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lhs_data : tensor or None
|
||||
The left operand, could be None if it's not required by op.
|
||||
rhs_data : tensor or None
|
||||
The right operand, could be None if it's not required by op.
|
||||
"""
|
||||
lhs_shape = F.shape(lhs_data)
|
||||
rhs_shape = F.shape(rhs_data)
|
||||
if len(lhs_shape) != len(rhs_shape):
|
||||
max_ndims = max(len(lhs_shape), len(rhs_shape))
|
||||
lhs_pad_ndims = max_ndims - len(lhs_shape)
|
||||
rhs_pad_ndims = max_ndims - len(rhs_shape)
|
||||
new_lhs_shape = (lhs_shape[0],) + (1,) * lhs_pad_ndims + lhs_shape[1:]
|
||||
new_rhs_shape = (rhs_shape[0],) + (1,) * rhs_pad_ndims + rhs_shape[1:]
|
||||
lhs_data = F.reshape(lhs_data, new_lhs_shape)
|
||||
rhs_data = F.reshape(rhs_data, new_rhs_shape)
|
||||
return lhs_data, rhs_data
|
||||
|
||||
|
||||
def gsddmm(g, op, lhs_data, rhs_data, lhs_target="u", rhs_target="v"):
|
||||
r"""Generalized Sampled-Dense-Dense Matrix Multiplication interface.
|
||||
It computes edge features by :attr:`op` lhs features and rhs features.
|
||||
|
||||
.. math::
|
||||
|
||||
x_{e} = \phi(x_{lhs}, x_{rhs}), \forall (u,e,v)\in \mathcal{G}
|
||||
|
||||
where :math:`x_{e}` is the returned feature on edges and :math:`x_u`,
|
||||
:math:`x_v` refers to :attr:`u`, :attr:`v` respectively. :math:`\phi`
|
||||
is the binary operator :attr:`op`, and :math:`\mathcal{G}` is the graph
|
||||
we apply gsddmm on: :attr:`g`. :math:`lhs` and :math:`rhs` are one of
|
||||
:math:`u,v,e`'s.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph.
|
||||
op : str
|
||||
Binary operator, could be ``add``, ``sub``, ``mul``, ``div``, ``dot``,
|
||||
``copy_lhs``, ``copy_rhs``.
|
||||
lhs_data : tensor or None
|
||||
The left operand, could be None if it's not required by op.
|
||||
rhs_data : tensor or None
|
||||
The right operand, could be None if it's not required by op.
|
||||
lhs_target: str
|
||||
Choice of ``u``(source), ``e``(edge) or ``v``(destination) for left operand.
|
||||
rhs_target: str
|
||||
Choice of ``u``(source), ``e``(edge) or ``v``(destination) for right operand.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
"""
|
||||
if g._graph.number_of_etypes() == 1:
|
||||
if op not in ["copy_lhs", "copy_rhs"]:
|
||||
lhs_data, rhs_data = reshape_lhs_rhs(lhs_data, rhs_data)
|
||||
return gsddmm_internal(
|
||||
g._graph, op, lhs_data, rhs_data, lhs_target, rhs_target
|
||||
)
|
||||
else:
|
||||
if op == "copy_lhs":
|
||||
rhs_data = [None] * g._graph.number_of_etypes()
|
||||
elif op == "copy_rhs":
|
||||
lhs_data = [None] * g._graph.number_of_ntypes()
|
||||
# TODO (Israt): Call reshape_lhs_rhs() on lhs and rhs data to match their dimension
|
||||
# and avoid broadcasting issue. Handle the case where different nodes have
|
||||
# different dimensions, and different etypes may need different broadcasting
|
||||
# dims for the same node.
|
||||
lhs_and_rhs_tuple = tuple(list(lhs_data) + list(rhs_data))
|
||||
return gsddmm_internal_hetero(
|
||||
g._graph,
|
||||
op,
|
||||
len(lhs_data),
|
||||
lhs_target,
|
||||
rhs_target,
|
||||
*lhs_and_rhs_tuple
|
||||
)
|
||||
|
||||
|
||||
def _gen_sddmm_func(lhs_target, rhs_target, binary_op):
|
||||
name = "{}_{}_{}".format(lhs_target, binary_op, rhs_target)
|
||||
target_dict = {"u": "source node", "e": "edge", "v": "destination node"}
|
||||
lhs_str = target_dict[lhs_target]
|
||||
rhs_str = target_dict[rhs_target]
|
||||
docstring = r"""Generalized SDDMM function.
|
||||
It computes edge features by {op} {lhs} features and {rhs} features.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph
|
||||
x : tensor
|
||||
The {lhs} features.
|
||||
y : tensor
|
||||
The {rhs} features.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function supports autograd (computing input gradients given the output gradient). If the
|
||||
feature shape of two input operands do not match, we first broadcasts the features to a unified
|
||||
shape (note that the memory usage will not increase accordingly) and then performs the operation.
|
||||
|
||||
Broadcasting follows NumPy semantics. Please see
|
||||
https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
|
||||
for more details about the NumPy broadcasting semantics.
|
||||
""".format(
|
||||
op=binary_op, lhs=lhs_str, rhs=rhs_str
|
||||
)
|
||||
|
||||
def func(g, x, y):
|
||||
return gsddmm(
|
||||
g, binary_op, x, y, lhs_target=lhs_target, rhs_target=rhs_target
|
||||
)
|
||||
|
||||
func.__name__ = name
|
||||
func.__doc__ = docstring
|
||||
return func
|
||||
|
||||
|
||||
def _register_sddmm_func():
|
||||
"""Register sddmm functions"""
|
||||
target = ["u", "v", "e"]
|
||||
for lhs, rhs in product(target, target):
|
||||
if lhs != rhs:
|
||||
for binary_op in ["add", "sub", "mul", "div", "dot"]:
|
||||
func = _gen_sddmm_func(lhs, rhs, binary_op)
|
||||
setattr(sys.modules[__name__], func.__name__, func)
|
||||
__all__.append(func.__name__)
|
||||
|
||||
|
||||
def copy_u(g, x):
|
||||
r"""Generalized SDDMM function that copies source node features to edges.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph.
|
||||
x : tensor
|
||||
The source node features.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function supports autograd (computing input gradients given the output gradient).
|
||||
"""
|
||||
return gsddmm(g, "copy_lhs", x, None)
|
||||
|
||||
|
||||
def copy_v(g, x):
|
||||
r"""Generalized SDDMM function that copies destination node features to edges.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph.
|
||||
x : tensor
|
||||
The destination node features.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function supports autograd (computing input gradients given the output gradient).
|
||||
"""
|
||||
return gsddmm(g, "copy_rhs", None, x)
|
||||
|
||||
|
||||
# pylint: disable=unused-argument
|
||||
def copy_e(g, x):
|
||||
r"""Generalized SDDMM function that copies destination node features to edges."""
|
||||
return x
|
||||
|
||||
|
||||
_register_sddmm_func()
|
||||
@@ -0,0 +1,130 @@
|
||||
"""Segment aggregation operators implemented using DGL graph."""
|
||||
|
||||
from .. import backend as F
|
||||
from ..base import DGLError
|
||||
|
||||
__all__ = ["segment_reduce", "segment_softmax", "segment_mm"]
|
||||
|
||||
|
||||
def segment_reduce(seglen, value, reducer="sum"):
|
||||
"""Segment reduction operator.
|
||||
|
||||
It aggregates the value tensor along the first dimension by segments.
|
||||
The first argument ``seglen`` stores the length of each segment. Its
|
||||
summation must be equal to the first dimension of the ``value`` tensor.
|
||||
Zero-length segments are allowed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
seglen : Tensor
|
||||
Segment lengths.
|
||||
value : Tensor
|
||||
Value to aggregate.
|
||||
reducer : str, optional
|
||||
Aggregation method. Can be 'sum', 'max', 'min', 'mean'.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor
|
||||
Aggregated tensor of shape ``(len(seglen), value.shape[1:])``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
>>> import dgl
|
||||
>>> import torch as th
|
||||
>>> val = th.ones(10, 3)
|
||||
>>> seg = th.tensor([1, 0, 5, 4]) # 4 segments
|
||||
>>> dgl.segment_reduce(seg, val)
|
||||
tensor([[1., 1., 1.],
|
||||
[0., 0., 0.],
|
||||
[5., 5., 5.],
|
||||
[4., 4., 4.]])
|
||||
"""
|
||||
offsets = F.cumsum(
|
||||
F.cat([F.zeros((1,), F.dtype(seglen), F.context(seglen)), seglen], 0), 0
|
||||
)
|
||||
if reducer == "mean":
|
||||
rst = F.segment_reduce("sum", value, offsets)
|
||||
rst_shape = F.shape(rst)
|
||||
z = F.astype(F.clamp(seglen, 1, len(value)), F.dtype(rst))
|
||||
z_shape = (rst_shape[0],) + (1,) * (len(rst_shape) - 1)
|
||||
return rst / F.reshape(z, z_shape)
|
||||
elif reducer in ["min", "sum", "max"]:
|
||||
rst = F.segment_reduce(reducer, value, offsets)
|
||||
if reducer in ["min", "max"]:
|
||||
rst = F.replace_inf_with_zero(rst)
|
||||
return rst
|
||||
else:
|
||||
raise DGLError("reducer {} not recognized.".format(reducer))
|
||||
|
||||
|
||||
def segment_softmax(seglen, value):
|
||||
"""Performa softmax on each segment.
|
||||
|
||||
The first argument ``seglen`` stores the length of each segment. Its
|
||||
summation must be equal to the first dimension of the ``value`` tensor.
|
||||
Zero-length segments are allowed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
seglen : Tensor
|
||||
Segment lengths.
|
||||
value : Tensor
|
||||
Value to aggregate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor
|
||||
Result tensor of the same shape as the ``value`` tensor.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
>>> import dgl
|
||||
>>> import torch as th
|
||||
>>> val = th.ones(10, 3)
|
||||
>>> seg = th.tensor([1, 0, 5, 4]) # 4 segments
|
||||
>>> dgl.segment_softmax(seg, val)
|
||||
tensor([[1.0000, 1.0000, 1.0000],
|
||||
[0.2000, 0.2000, 0.2000],
|
||||
[0.2000, 0.2000, 0.2000],
|
||||
[0.2000, 0.2000, 0.2000],
|
||||
[0.2000, 0.2000, 0.2000],
|
||||
[0.2000, 0.2000, 0.2000],
|
||||
[0.2500, 0.2500, 0.2500],
|
||||
[0.2500, 0.2500, 0.2500],
|
||||
[0.2500, 0.2500, 0.2500],
|
||||
[0.2500, 0.2500, 0.2500]])
|
||||
"""
|
||||
value_max = segment_reduce(seglen, value, reducer="max")
|
||||
value = F.exp(value - F.repeat(value_max, seglen, dim=0))
|
||||
value_sum = segment_reduce(seglen, value, reducer="sum")
|
||||
return value / F.repeat(value_sum, seglen, dim=0)
|
||||
|
||||
|
||||
def segment_mm(a, b, seglen_a):
|
||||
r"""Performs matrix multiplication according to segments.
|
||||
|
||||
Suppose ``seglen_a == [10, 5, 0, 3]``, the operator will perform
|
||||
four matrix multiplications::
|
||||
|
||||
a[0:10] @ b[0], a[10:15] @ b[1],
|
||||
a[15:15] @ b[2], a[15:18] @ b[3]
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a : Tensor
|
||||
The left operand, 2-D tensor of shape ``(N, D1)``
|
||||
b : Tensor
|
||||
The right operand, 3-D tensor of shape ``(R, D1, D2)``
|
||||
seglen_a : Tensor
|
||||
An integer tensor of shape ``(R,)``. Each element is the length of segments
|
||||
of input ``a``. The summation of all elements must be equal to ``N``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor
|
||||
The output dense matrix of shape ``(N, D2)``
|
||||
"""
|
||||
return F.segment_mm(a, b, seglen_a)
|
||||
@@ -0,0 +1,241 @@
|
||||
"""Internal module for general spmm operators."""
|
||||
import sys
|
||||
|
||||
from .. import backend as F
|
||||
from ..backend import (
|
||||
gspmm as gspmm_internal,
|
||||
gspmm_hetero as gspmm_internal_hetero,
|
||||
)
|
||||
|
||||
__all__ = ["gspmm"]
|
||||
|
||||
|
||||
def reshape_lhs_rhs(lhs_data, rhs_data):
|
||||
r"""Expand dims so that there will be no broadcasting issues with different
|
||||
number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
|
||||
that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
|
||||
(E, 5, 3, 4)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lhs_data : tensor or None
|
||||
The left operand, could be None if it's not required by op.
|
||||
rhs_data : tensor or None
|
||||
The right operand, could be None if it's not required by op.
|
||||
"""
|
||||
lhs_shape = F.shape(lhs_data)
|
||||
rhs_shape = F.shape(rhs_data)
|
||||
if len(lhs_shape) != len(rhs_shape):
|
||||
max_ndims = max(len(lhs_shape), len(rhs_shape))
|
||||
lhs_pad_ndims = max_ndims - len(lhs_shape)
|
||||
rhs_pad_ndims = max_ndims - len(rhs_shape)
|
||||
new_lhs_shape = (lhs_shape[0],) + (1,) * lhs_pad_ndims + lhs_shape[1:]
|
||||
new_rhs_shape = (rhs_shape[0],) + (1,) * rhs_pad_ndims + rhs_shape[1:]
|
||||
lhs_data = F.reshape(lhs_data, new_lhs_shape)
|
||||
rhs_data = F.reshape(rhs_data, new_rhs_shape)
|
||||
return lhs_data, rhs_data
|
||||
|
||||
|
||||
def gspmm(g, op, reduce_op, lhs_data, rhs_data):
|
||||
r"""Generalized Sparse Matrix Multiplication interface.
|
||||
It fuses two steps into one kernel.
|
||||
|
||||
1. Computes messages by :attr:`op` source node and edge features.
|
||||
2. Aggregate the messages by :attr:`reduce_op` as the features on destination nodes.
|
||||
|
||||
.. math::
|
||||
x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))
|
||||
|
||||
where :math:`x_v` is the returned feature on destination nodes, and :math:`x_u`,
|
||||
:math:`x_e` refers to :attr:`u`, :attr:`e` respectively. :math:`\rho` means binary
|
||||
operator :attr:`op` and :math:`\psi` means reduce operator :attr:`reduce_op`,
|
||||
:math:`\mathcal{G}` is the graph we apply gspmm on: :attr:`g`.
|
||||
|
||||
Note that this function does not handle gradients.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph.
|
||||
op : str
|
||||
The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
|
||||
``copy_lhs``, ``copy_rhs``.
|
||||
reduce_op : str
|
||||
Reduce operator, could be ``sum``, ``max``, ``min``, ``mean``.
|
||||
lhs_data : tensor or None
|
||||
The left operand, could be None if it's not required by the op.
|
||||
rhs_data : tensor or None
|
||||
The right operand, could be None if it's not required by the op.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
"""
|
||||
if g._graph.number_of_etypes() == 1:
|
||||
if op not in ["copy_lhs", "copy_rhs"]:
|
||||
lhs_data, rhs_data = reshape_lhs_rhs(lhs_data, rhs_data)
|
||||
# With max and min reducers infinity will be returned for zero degree nodes
|
||||
ret = gspmm_internal(
|
||||
g._graph,
|
||||
op,
|
||||
"sum" if reduce_op == "mean" else reduce_op,
|
||||
lhs_data,
|
||||
rhs_data,
|
||||
)
|
||||
else:
|
||||
# lhs_data or rhs_data is None only in unary functions like ``copy-u`` or ``copy_e``
|
||||
lhs_data = (
|
||||
[None] * g._graph.number_of_ntypes()
|
||||
if lhs_data is None
|
||||
else lhs_data
|
||||
)
|
||||
rhs_data = (
|
||||
[None] * g._graph.number_of_etypes()
|
||||
if rhs_data is None
|
||||
else rhs_data
|
||||
)
|
||||
# TODO (Israt): Call reshape func
|
||||
lhs_and_rhs_tuple = tuple(list(lhs_data) + list(rhs_data))
|
||||
ret = gspmm_internal_hetero(
|
||||
g._graph,
|
||||
op,
|
||||
"sum" if reduce_op == "mean" else reduce_op,
|
||||
len(lhs_data),
|
||||
*lhs_and_rhs_tuple
|
||||
)
|
||||
# TODO (Israt): Add support for 'mean' in heterograph
|
||||
# divide in degrees for mean reducer.
|
||||
if reduce_op == "mean":
|
||||
ret_shape = F.shape(ret)
|
||||
deg = g.in_degrees()
|
||||
deg = F.astype(F.clamp(deg, 1, max(g.num_edges(), 1)), F.dtype(ret))
|
||||
deg_shape = (ret_shape[0],) + (1,) * (len(ret_shape) - 1)
|
||||
return ret / F.reshape(deg, deg_shape)
|
||||
else:
|
||||
return ret
|
||||
|
||||
|
||||
def _attach_zerodeg_note(docstring, reducer):
|
||||
note1 = """
|
||||
The {} function will return zero for nodes with no incoming messages.""".format(
|
||||
reducer
|
||||
)
|
||||
note2 = """
|
||||
This is implemented by replacing all {} values to zero.
|
||||
""".format(
|
||||
"infinity" if reducer == "min" else "negative infinity"
|
||||
)
|
||||
|
||||
docstring = docstring + note1
|
||||
if reducer in ("min", "max"):
|
||||
docstring = docstring + note2
|
||||
return docstring
|
||||
|
||||
|
||||
def _gen_spmm_func(binary_op, reduce_op):
|
||||
name = "u_{}_e_{}".format(binary_op, reduce_op)
|
||||
docstring = """Generalized SpMM function.
|
||||
It fuses two steps into one kernel.
|
||||
|
||||
1. Computes messages by {} source node and edge features.
|
||||
2. Aggregate the messages by {} as the features on destination nodes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph
|
||||
x : tensor
|
||||
The source node features.
|
||||
y : tensor
|
||||
The edge features.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function supports autograd (computing input gradients given the output gradient). If the
|
||||
feature shape of two input operands do not match, we first broadcasts the features to a unified
|
||||
shape (note that the memory usage will not increase accordingly) and then performs the operation.
|
||||
|
||||
Broadcasting follows NumPy semantics. Please see
|
||||
https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
|
||||
for more details about the NumPy broadcasting semantics.
|
||||
""".format(
|
||||
binary_op, reduce_op
|
||||
)
|
||||
docstring = _attach_zerodeg_note(docstring, reduce_op)
|
||||
|
||||
def func(g, x, y):
|
||||
return gspmm(g, binary_op, reduce_op, x, y)
|
||||
|
||||
func.__name__ = name
|
||||
func.__doc__ = docstring
|
||||
return func
|
||||
|
||||
|
||||
def _gen_copy_reduce_func(binary_op, reduce_op):
|
||||
|
||||
name = "{}_{}".format(binary_op, reduce_op)
|
||||
binary_str = {
|
||||
"copy_u": "It copies node feature to edge as the message.",
|
||||
"copy_e": "It regards edge feature as message.",
|
||||
}
|
||||
x_str = {"copy_u": "source node", "copy_e": "edge"}
|
||||
docstring = lambda binary_op: _attach_zerodeg_note(
|
||||
"""Generalized SpMM function. {}
|
||||
Then aggregates the message by {} on destination nodes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The input graph
|
||||
x : tensor
|
||||
The {} features.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor
|
||||
The result tensor.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function supports autograd (computing input gradients given the output gradient).
|
||||
""".format(
|
||||
binary_str[binary_op], reduce_op, x_str[binary_op]
|
||||
),
|
||||
reduce_op,
|
||||
)
|
||||
|
||||
def func(g, x):
|
||||
if binary_op == "copy_u":
|
||||
return gspmm(g, "copy_lhs", reduce_op, x, None)
|
||||
else:
|
||||
return gspmm(g, "copy_rhs", reduce_op, None, x)
|
||||
|
||||
func.__name__ = name
|
||||
func.__doc__ = docstring(binary_op)
|
||||
return func
|
||||
|
||||
|
||||
def _register_spmm_func():
|
||||
"""Register spmm functions
|
||||
|
||||
- Binary operation plus reduction between u and e: u_[]_e_[]
|
||||
- Copy u plus reduction: copy_u_[]
|
||||
- Copy e plus reduction: copy_e_[]
|
||||
"""
|
||||
for binary_op in ["add", "sub", "mul", "div", "copy_u", "copy_e"]:
|
||||
for reduce_op in ["sum", "max", "min", "mean"]:
|
||||
if binary_op.startswith("copy"):
|
||||
func = _gen_copy_reduce_func(binary_op, reduce_op)
|
||||
else:
|
||||
func = _gen_spmm_func(binary_op, reduce_op)
|
||||
setattr(sys.modules[__name__], func.__name__, func)
|
||||
__all__.append(func.__name__)
|
||||
|
||||
|
||||
_register_spmm_func()
|
||||
Reference in New Issue
Block a user