dgl.ops

用于图上消息传递的框架无关操作符。

GSpMM 函数

广义稀疏矩阵-稠密矩阵乘法函数。 它将融合两个步骤到一个内核中。

1. 通过加/减/乘/除源节点和边特征来计算消息，或将节点特征复制到边。

2. 通过求和/最大值/最小值/平均值将消息聚合为目标节点的特征。

我们的实现支持在PyTorch/MXNet/Tensorflow中的CPU/GPU上的张量作为输入。所有操作符都配备了自动求导（计算给定输出梯度的输入梯度）和广播（如果操作数的特征形状不匹配，我们首先将它们广播到相同的形状，然后应用二元操作符）。我们的广播语义遵循NumPy，请参阅 https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html 了解更多详情。

我们所说的fuses是指消息不在边上具体化，而是直接在目标节点上计算结果，从而节省内存成本。GSpMM操作的空间复杂度是\(O(|N|D)\)，其中\(|N|\)指的是图中节点的数量，\(D\)指的是特征的大小（如果你的特征是一个多维张量，则\(D=\prod_{i=1}^{N}D_i\)）。

以下是一个展示GSpMM如何工作的示例（我们在这里使用PyTorch作为后端，您可以通过类似的使用方式在其他框架上享受相同的便利）：

>>> import dgl
>>> import torch as th
>>> import dgl.ops as F
>>> g = dgl.graph(([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]))  # 3 nodes, 6 edges
>>> x = th.ones(3, 2, requires_grad=True)
>>> x
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]], requires_grad=True)
>>> y = th.arange(1, 13).float().view(6, 2).requires_grad_()
tensor([[ 1.,  2.],
        [ 3.,  4.],
        [ 5.,  6.],
        [ 7.,  8.],
        [ 9., 10.],
        [11., 12.]], requires_grad=True)
>>> out_1 = F.u_mul_e_sum(g, x, y)
>>> out_1  # (10, 12) = ((1, 1) * (3, 4)) + ((1, 1) * (7, 8))
tensor([[ 1.,  2.],
        [10., 12.],
        [25., 28.]], grad_fn=<GSpMMBackward>)
>>> out_1.sum().backward()
>>> x.grad
tensor([[12., 15.],
        [18., 20.],
        [12., 13.]])
>>> y.grad
tensor([[1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.]])
>>> out_2 = F.copy_u_sum(g, x)
>>> out_2
tensor([[1., 1.],
        [2., 2.],
        [3., 3.]], grad_fn=<GSpMMBackward>)
>>> out_3 = F.u_add_e_max(g, x, y)
>>> out_3
tensor([[ 2.,  3.],
        [ 8.,  9.],
        [12., 13.]], grad_fn=<GSpMMBackward>)
>>> y1 = th.rand(6, 4, 2, requires_grad=True)  # test broadcast
>>> F.u_mul_e_sum(g, x, y1).shape  # (2,), (4, 2) -> (4, 2)
torch.Size([3, 4, 2])

对于所有操作符，输入图可以是同构图或二分图。

gspmm(g, op, reduce_op, lhs_data, rhs_data)	广义稀疏矩阵乘法接口。
u_add_e_sum(g, x, y)	广义SpMM函数。
u_sub_e_sum(g, x, y)	广义SpMM函数。
u_mul_e_sum(g, x, y)	广义SpMM函数。
u_div_e_sum(g, x, y)	广义SpMM函数。
u_add_e_max(g, x, y)	广义SpMM函数。
u_sub_e_max(g, x, y)	广义SpMM函数。
u_mul_e_max(g, x, y)	广义SpMM函数。
u_div_e_max(g, x, y)	广义SpMM函数。
u_add_e_min(g, x, y)	广义SpMM函数。
u_sub_e_min(g, x, y)	广义SpMM函数。
u_mul_e_min(g, x, y)	广义SpMM函数。
u_div_e_min(g, x, y)	广义SpMM函数。
u_add_e_mean(g, x, y)	广义SpMM函数。
u_sub_e_mean(g, x, y)	广义SpMM函数。
u_mul_e_mean(g, x, y)	广义SpMM函数。
u_div_e_mean(g, x, y)	广义SpMM函数。
copy_u_sum(g, x)	广义SpMM函数。
copy_e_sum(g, x)	广义SpMM函数。
copy_u_max(g, x)	广义SpMM函数。
copy_e_max(g, x)	广义SpMM函数。
copy_u_min(g, x)	广义SpMM函数。
copy_e_min(g, x)	广义SpMM函数。
copy_u_mean(g, x)	广义SpMM函数。
copy_e_mean(g, x)	广义SpMM函数。

GSDDMM 函数

广义采样密集-密集矩阵乘法。 它通过在源/目标节点或边上进行加/减/乘/除/点积特征来计算边特征。

与GSpMM类似，我们的实现支持PyTorch/MXNet/Tensorflow中的CPU/GPU张量作为输入。所有操作符都配备了自动求导和广播功能。

GSDDMM的内存成本是\(O(|E|D)\)，其中\(|E|\)指的是图中的边数，而\(D\)指的是特征大小。

请注意，我们支持dot操作符，其语义上与通过求和减少最后一个维度的mul操作符的结果相同。然而，dot操作符在内存使用上更为高效，因为它融合了mul和求和减少操作，这在最后一个维度上的特征大小不可忽视的情况下（例如，类似Transformer模型中的多头注意力机制）尤为重要。

以下是一个展示GSDDMM如何工作的示例：

>>> import dgl
>>> import torch as th
>>> import dgl.ops as F
>>> g = dgl.graph(([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]))  # 3 nodes, 6 edges
>>> x = th.ones(3, 2, requires_grad=True)
>>> x
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]], requires_grad=True)
>>> y = th.arange(1, 7).float().view(3, 2).requires_grad_()
>>> y
tensor([[1., 2.],
        [3., 4.],
        [5., 6.]], requires_grad=True)
>>> e = th.ones(6, 1, 2, requires_grad=True) * 2
tensor([[[2., 2.]],
        [[2., 2.]],
        [[2., 2.]],
        [[2., 2.]],
        [[2., 2.]],
        [[2., 2.]]], grad_fn=<MulBackward0>)
>>> out1 = F.u_div_v(g, x, y)
tensor([[1.0000, 0.5000],
        [0.3333, 0.2500],
        [0.2000, 0.1667],
        [0.3333, 0.2500],
        [0.2000, 0.1667],
        [0.2000, 0.1667]], grad_fn=<GSDDMMBackward>)
>>> out1.sum().backward()
>>> x.grad
tensor([[1.5333, 0.9167],
        [0.5333, 0.4167],
        [0.2000, 0.1667]])
>>> y.grad
tensor([[-1.0000, -0.2500],
        [-0.2222, -0.1250],
        [-0.1200, -0.0833]])
>>> out2 = F.e_sub_v(g, e, y)
>>> out2
tensor([[[ 1.,  0.]],
        [[-1., -2.]],
        [[-3., -4.]],
        [[-1., -2.]],
        [[-3., -4.]],
        [[-3., -4.]]], grad_fn=<GSDDMMBackward>)
>>> out3 = F.copy_v(g, y)
>>> out3
tensor([[1., 2.],
        [3., 4.],
        [5., 6.],
        [3., 4.],
        [5., 6.],
        [5., 6.]], grad_fn=<GSDDMMBackward>)
>>> out4 = F.u_dot_v(g, x, y)
>>> out4  # the last dimension was reduced to size 1.
tensor([[ 3.],
        [ 7.],
        [11.],
        [ 7.],
        [11.],
        [11.]], grad_fn=<GSDDMMBackward>)

gsddmm(g, op, lhs_data, rhs_data[, ...])	广义采样-密集-密集矩阵乘法接口。
u_add_v(g, x, y)	广义SDDMM函数。
u_sub_v(g, x, y)	广义SDDMM函数。
u_mul_v(g, x, y)	广义SDDMM函数。
u_dot_v(g, x, y)	广义SDDMM函数。
u_div_v(g, x, y)	广义SDDMM函数。
u_add_e(g, x, y)	广义SDDMM函数。
u_sub_e(g, x, y)	广义SDDMM函数。
u_mul_e(g, x, y)	广义SDDMM函数。
u_dot_e(g, x, y)	广义SDDMM函数。
u_div_e(g, x, y)	广义SDDMM函数。
e_add_v(g, x, y)	广义SDDMM函数。
e_sub_v(g, x, y)	广义SDDMM函数。
e_mul_v(g, x, y)	广义SDDMM函数。
e_dot_v(g, x, y)	广义SDDMM函数。
e_div_v(g, x, y)	广义SDDMM函数。
v_add_u(g, x, y)	广义SDDMM函数。
v_sub_u(g, x, y)	广义SDDMM函数。
v_mul_u(g, x, y)	广义SDDMM函数。
v_dot_u(g, x, y)	广义SDDMM函数。
v_div_u(g, x, y)	广义SDDMM函数。
e_add_u(g, x, y)	广义SDDMM函数。
e_sub_u(g, x, y)	广义SDDMM函数。
e_mul_u(g, x, y)	广义SDDMM函数。
e_dot_u(g, x, y)	广义SDDMM函数。
e_div_u(g, x, y)	广义SDDMM函数。
v_add_e(g, x, y)	广义SDDMM函数。
v_sub_e(g, x, y)	广义SDDMM函数。
v_mul_e(g, x, y)	广义SDDMM函数。
v_dot_e(g, x, y)	广义SDDMM函数。
v_div_e(g, x, y)	广义SDDMM函数。
copy_u(g, x)	广义SDDMM函数，将源节点特征复制到边。
copy_v(g, x)	广义SDDMM函数，将目标节点特征复制到边。

与GSpMM类似，GSDDMM操作符支持同质图和二分图。

分段减少模块

DGL 提供了操作符来按段沿第一维度减少值张量。

segment_reduce(seglen, value[, reducer])

分段归约操作符。

GatherMM 和 SegmentMM 模块

SegmentMM: DGL 提供了根据段执行矩阵乘法的操作符。

GatherMM: DGL 提供了根据给定索引收集数据并执行矩阵乘法的操作符。

gather_mm(a, b, *, idx_b)	根据给定的索引收集数据并执行矩阵乘法。
segment_mm(a, b, seglen_a)	根据段执行矩阵乘法。

支持的数据类型

在dgl.ops中定义的运算符支持浮点数据类型，即操作数必须是half (float16) /float/double张量。 输入张量必须具有相同的数据类型（如果一个输入张量的类型是float16，而另一个输入张量的数据类型是float32，用户必须将其中一个转换为与另一个对齐）。

float16 数据类型支持默认是禁用的，因为它对GPU的计算能力有最低要求，即 sm_53（Pascal、Volta、Turing和Ampere架构）。

User can enable float16 for mixed precision training by compiling DGL from source (see Mixed Precision Training tutorial for details).

与消息传递API的关系

dgl.update_all 和 dgl.apply_edges 调用内置的消息/减少函数 将被分派到 dgl.ops 中定义的操作符函数调用：

>>> import dgl
>>> import torch as th
>>> import dgl.ops as F
>>> import dgl.function as fn
>>> g = dgl.rand_graph(100, 1000)   # create a DGLGraph with 100 nodes and 1000 edges.
>>> x = th.rand(100, 20)            # node features.
>>> e = th.rand(1000, 20)
>>>
>>> # dgl.update_all + builtin functions
>>> g.srcdata['x'] = x              # srcdata is the same as ndata for graphs with one node type.
>>> g.edata['e'] = e
>>> g.update_all(fn.u_mul_e('x', 'e', 'm'), fn.sum('m', 'y'))
>>> y = g.dstdata['y']              # dstdata is the same as ndata for graphs with one node type.
>>>
>>> # use GSpMM operators defined in dgl.ops directly
>>> y = F.u_mul_e_sum(g, x, e)
>>>
>>> # dgl.apply_edges + builtin functions
>>> g.srcdata['x'] = x
>>> g.dstdata['y'] = y
>>> g.apply_edges(fn.u_dot_v('x', 'y', 'z'))
>>> z = g.edata['z']
>>>
>>> # use GSDDMM operators defined in dgl.ops directly
>>> z = F.u_dot_v(g, x, y)

由用户决定是否使用消息传递API或GSpMM/GSDDMM操作符，两者的效率相同。使用消息传递API编写的程序看起来更像DGL风格，但在某些情况下调用GSpMM/GSDDMM操作符更为简洁。

请注意，在PyTorch中，所有在dgl.ops中定义的运算符都支持高阶梯度，因此消息传递API也支持，因为它们完全依赖于这些运算符。