Given x y z index, change to the linear index. (matlab's sub2ind) Args: idx: (*, n, 3) long size: (*, 3) long Returns: (*, n) long
(
idx: torch.Tensor, # (*, n, 3)
size: torch.Tensor, # (*, 3)
)
| 12 | |
| 13 | |
| 14 | def sub2ind( |
| 15 | idx: torch.Tensor, # (*, n, 3) |
| 16 | size: torch.Tensor, # (*, 3) |
| 17 | ) -> torch.Tensor: |
| 18 | """ |
| 19 | Given x y z index, change to the linear index. (matlab's sub2ind) |
| 20 | |
| 21 | Args: |
| 22 | idx: |
| 23 | (*, n, 3) long |
| 24 | size: |
| 25 | (*, 3) long |
| 26 | |
| 27 | Returns: |
| 28 | (*, n) long |
| 29 | """ |
| 30 | # linear_idx = idx[..., 0] + \ |
| 31 | # idx[..., 1] * size[..., 0:1] + \ |
| 32 | # idx[..., 2] * (size[..., 0:1] * size[..., 1:2]) # (*, n) |
| 33 | |
| 34 | linear_idx = idx[..., 2] + \ |
| 35 | idx[..., 1] * size[..., 2:3] + \ |
| 36 | idx[..., 0] * (size[..., 1:2] * size[..., 2:3]) # (*, n) |
| 37 | |
| 38 | return linear_idx |
| 39 | |
| 40 | |
| 41 | def ind2sub( |
no outgoing calls
no test coverage detected