(pos: torch.Tensor)
| 64 | |
| 65 | def sort_splats_morton(splats: Dict[str, Tensor], verbose: bool = True, return_indices: bool = False) -> Dict[str, Tensor]: |
| 66 | def mortonEncode(pos: torch.Tensor) -> torch.Tensor: |
| 67 | def splitBy3(a): |
| 68 | x = a & 0x1FFFFF # we only look at the first 21 bits |
| 69 | x = (x | x << 32) & 0x1F00000000FFFF |
| 70 | x = (x | x << 16) & 0x1F0000FF0000FF |
| 71 | x = (x | x << 8) & 0x100F00F00F00F00F |
| 72 | x = (x | x << 4) & 0x10C30C30C30C30C3 |
| 73 | x = (x | x << 2) & 0x1249249249249249 |
| 74 | return x |
| 75 | x, y, z = pos.unbind(-1) |
| 76 | answer = torch.zeros(len(pos), dtype=torch.long, device=pos.device) |
| 77 | answer |= splitBy3(x) | splitBy3(y) << 1 | splitBy3(z) << 2 |
| 78 | return answer |
| 79 | |
| 80 | with torch.no_grad(): |
| 81 | xyz_q = ( |
no test coverage detected