()
| 71 | |
| 72 | |
| 73 | def test_argsort_2d(): |
| 74 | M, N = 4, 128 # noqa: N806 |
| 75 | x = torch.randn(M, N, dtype=torch.float32, device="cuda") |
| 76 | out_vals = torch.empty_like(x) |
| 77 | out_ids = torch.empty(M, N, dtype=torch.int32, device="cuda") |
| 78 | |
| 79 | argsort_2d_kernel[(M,)](x, out_vals, out_ids, stride_m=N, N=N, M=M) |
| 80 | |
| 81 | expected_vals, expected_ids = torch.sort(x, dim=-1, descending=True) |
| 82 | assert torch.allclose(out_vals, expected_vals) |
| 83 | assert torch.equal(out_ids, expected_ids.to(torch.int32)) |
nothing calls this directly
no outgoing calls
no test coverage detected