edge_softmax

class edge_softmax(graph: GraphShardManager, logits: Tensor, eids='__ALL__', norm_by: str = 'dst')

Implements a similar functionality as DGL’s dgl.nn.edge_softmax on distributed graphs.

Only supports a subset of the possible argument values.

Parameters:
  • graph (GraphShardManager) – The distributed graph

  • logits (torch.Tensor) – The edge logits. The size of the first dimension should be the same as the number of edges in the graph argument

  • eids – must be dgl.ALL

  • norm_by (str) – must be 'dst'

Returns:

A tensor with the same size as logits contaning the softmax-normalized logits