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.

xv=ψ(u,v,e)∈G(ρ(xu,xe))

where xv is the returned feature on destination nodes, and xu, xe refers to u, e respectively. ρ means binary operator op and ψ means reduce operator reduce_op, 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