| 8 | |
| 9 | class TestGetGridIdx(unittest.TestCase): |
| 10 | def _test( |
| 11 | self, |
| 12 | b: int, |
| 13 | n: int, |
| 14 | grid_size, |
| 15 | center, |
| 16 | grid_width, |
| 17 | ): |
| 18 | points = torch.rand(b, n, 3) |
| 19 | gidx, valid_mask = naive.get_grid_idx( |
| 20 | points=points, |
| 21 | grid_size=grid_size, |
| 22 | center=center, |
| 23 | grid_width=grid_width, |
| 24 | ) |
| 25 | if isinstance(grid_size, int): |
| 26 | total_cells = grid_size * grid_size * grid_size |
| 27 | else: |
| 28 | total_cells = torch.prod(grid_size, dim=-1, keepdim=True) # (*, 1) |
| 29 | torch.prod(gidx) |
| 30 | assert torch.logical_or(gidx >= 0, torch.logical_not(valid_mask)).all() |
| 31 | assert torch.logical_or(gidx < total_cells, torch.logical_not(valid_mask)).all() |
| 32 | |
| 33 | |
| 34 | def test1(self): |