dgl.ops.gspmm

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

Generalized Sparse Matrix Multiplication interface. It fuses two steps into one kernel.

  1. Computes messages by op source node and edge features.

  2. Aggregate the messages by reduce_op as the features on destination nodes.

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

where \(x_v\) is the returned feature on destination nodes, and \(x_u\), \(x_e\) refers to u, e respectively. \(\rho\) means binary operator op and \(\psi\) means reduce operator reduce_op, \(\mathcal{G}\) is the graph we apply gspmm on: 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

The result tensor.

Return type

tensor