(self, b=10, n=20, max_size=20)
| 17 | |
| 18 | class TestIndexing(unittest.TestCase): |
| 19 | def test_sub2ind(self, b=10, n=20, max_size=20): |
| 20 | size = torch.randint(max_size-1, size=[b, 3]) + 1 # (b, 3) |
| 21 | total_cells = torch.prod(size, dim=-1) # (b,) |
| 22 | subidx = torch.rand(b, n, 3).floor().long() |
| 23 | |
| 24 | # python |
| 25 | stime = timer() |
| 26 | ind_gt = naive.sub2ind(subidx, size) |
| 27 | time_python = timer() - stime |
| 28 | |
| 29 | # cpp |
| 30 | stime = timer() |
| 31 | ind_cpp = pr_cpp.sub2ind(subidx, size) |
| 32 | time_cpp = timer() - stime |
| 33 | |
| 34 | print('sub2ind') |
| 35 | print(f'python: {time_python * 1000.:.3f} ms') |
| 36 | print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python/time_cpp:.2f} speed up)') |
| 37 | |
| 38 | assert torch.allclose(ind_gt, ind_cpp) |
| 39 | |
| 40 | |
| 41 | def test_ind2sub(self, b=10, n=20, max_size=20): |
nothing calls this directly
no outgoing calls
no test coverage detected