gammagl.utils.segment_softmax

class segment_softmax(data, segment_ids, num_segments)[source]

segment softmax function.

Parameters:
  • data – The source tensor.

  • segment_ids – The indices of elements for applying the softmax.

  • num_segments – The number of segments.

Returns:

softmax score.

Return type:

tensor