(N)
| 59 | |
| 60 | @pytest.mark.parametrize("N", [16, 64, 128, 256]) |
| 61 | def test_argsort_ascending(N): # noqa: N803 |
| 62 | x = torch.randn(N, dtype=torch.float32, device="cuda") |
| 63 | out_vals = torch.empty_like(x) |
| 64 | out_ids = torch.empty(N, dtype=torch.int32, device="cuda") |
| 65 | |
| 66 | argsort_asc_kernel[(1,)](x, out_vals, out_ids, N=N) |
| 67 | |
| 68 | expected_vals, expected_ids = torch.sort(x, descending=False) |
| 69 | assert torch.allclose(out_vals, expected_vals) |
| 70 | assert torch.equal(out_ids, expected_ids.to(torch.int32)) |
| 71 | |
| 72 | |
| 73 | def test_argsort_2d(): |
nothing calls this directly
no outgoing calls
no test coverage detected