class_docstr: MathTest
      
docstr: pgl.math.segment_topk test
          
docstr_tok: ["pgl.math.segment_topk","test"]
code: def test_segment_topk(self):
      """pgl.math.segment_topk test"""
      data = paddle.to_tensor(
          [[1, 2, 3], [3, 2, 1], [4, 5, 6], [9, 9, 8], [20, 1, 5]],
          dtype="float32",
      )
      segment_ids = paddle.to_tensor([0, 0, 1, 1, 1], dtype="int64")
      scores = paddle.to_tensor([1, 3, 2, 7, 4], dtype="float32")
      output, index = pgl.math.segment_topk(
          data, scores, segment_ids, 0.5, return_index=True
      )
      np.testing.assert_almost_equal(index.numpy().tolist(), [1, 3, 4])
      np.testing.assert_almost_equal(
          output.numpy().tolist(), [[3, 2, 1], [9, 9, 8], [20, 1, 5]]
      )
      output, index = pgl.math.segment_topk(
          data,
          scores,
          segment_ids,
          0.5,
          min_score=2.1,
          return_index=True,
      )
      np.testing.assert_almost_equal(index.numpy().tolist(), [1, 3, 4])
  
file_path: PGL/tests/test_math.py
imports: import os
  import sys
  import re
  import time
  import unittest
  import numpy
  import paddle
  import pgl
  from pgl.utils.logger import log
label_desc: Use image node_id 3 for calling the MathTest obj's underlying member method code with example usage: obj.test_segment_topk() without return types