| 671 | |
| 672 | |
| 673 | class TestMaxHeap(unittest.TestCase): |
| 674 | def test(self, b=1, m=1000, n=1000, seed=0): |
| 675 | if seed is not None: |
| 676 | torch.manual_seed(seed) |
| 677 | random.seed(seed) |
| 678 | np.random.seed(seed) |
| 679 | values = torch.rand(b, m, n).cuda() |
| 680 | k = 40 |
| 681 | print(values.device) |
| 682 | min_k_values = pr_cuda.keep_min_k_values( |
| 683 | values, |
| 684 | k |
| 685 | ) |
| 686 | min_k_values_sorted = min_k_values.sort(dim=-1)[0] |
| 687 | min_k_values_sorted_gt = values.sort(dim=-1)[0][..., :k] |
| 688 | assert (min_k_values_sorted == min_k_values_sorted_gt).all() |
| 689 | |
| 690 | |
| 691 | if __name__ == '__main__': |
nothing calls this directly
no outgoing calls
no test coverage detected