code: def segment_softmax(data, segment_ids):
      """
      Segment softmax operator.
  
      This operator calculate the softmax elements of input `data` which with
      the same index in `segment_ids`.
  
      Args:
          data (tensor): a tensor, available data type float32, float64.
          segment_ids (tensor): a 1-d tensor, which have the same size
                              with the first dimension of input data.
                              available data type is int32, int64.
  
      Returns:
         output (Tensor): the softmax result.
  
      Examples:
  
          .. code-block:: python
  
              import paddle
              import pgl
              data = [[1, 2, 3],
                      [3, 2, 1],
                      [4, 5, 6]]
              data = paddle.to_tensor(, dtype='float32')
              segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
              out = pgl.math.segment_softmax(data, segment_ids)
  
              # Outputs:
                      [[0.11920292 0.5        0.880797  ]
                       [0.880797   0.5        0.11920292]
                       [1.         1.         1.        ]]
  
      """
      with paddle.no_grad():
          # no need gradients
          data_max = segment_max(data, segment_ids)
          data_max = paddle.gather(data_max, segment_ids, axis=0)
      data = data - data_max
      data = paddle.exp(data)
      sum_data = segment_sum(data, segment_ids)
      sum_data = paddle.gather(sum_data, segment_ids, axis=0)
      return data / sum_data
  
file_path: PGL/pgl/math.py
imports: import paddle
  from pgl.utils.op import get_index_from_counts
label_desc: Use image node_id 6 for calling a global function with example usage: segment_softmax(data, segment_ids) and returns: unknown