dgl.ops.gspmm

dgl.ops.gspmm(g, op, reduce_op, lhs_data, rhs_data)

广义稀疏矩阵乘法接口。它将两个步骤融合到一个核中。

1. 通过 op 计算源节点和边特征的消息。

2. 通过 reduce_op 聚合消息作为目标节点上的特征。

\[x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))\]

其中 \(x_v\) 是目标节点上的返回特征，\(x_u\), \(x_e\) 分别指代 u, e。 \(\rho\) 表示二元操作符 op，\(\psi\) 表示归约操作符 reduce_op，\(\mathcal{G}\) 是我们应用 gspmm 的图：g。

注意：此函数不处理梯度。

参数:

• g (DGLGraph) – 输入图。

• op (str) – 二元操作符的名称，可以是 add, sub, mul, div, copy_lhs, copy_rhs。

• reduce_op (str) – 归约操作符，可以是 sum, max, min, mean。

• lhs_data (tensor 或 None) – 左操作数，如果 op 不需要，可以为 None。

• rhs_data (tensor 或 None) – 右操作数，如果 op 不需要，可以为 None。

返回值:

结果张量。

返回类型:

tensor