Given linear index, return (i,j,k) Args: ind: (*, n) int64_t, z + y * sz + x * sy * sz size: (*, 3) int64_t Returns: (*, n, 3) int64_t
| 127 | // (*, n, 3) int64_t |
| 128 | // |
| 129 | torch::Tensor ind2sub( |
| 130 | const torch::Tensor & ind, // (*, n) int64_t |
| 131 | const torch::Tensor & size // (*, 3) int64_t |
| 132 | ) { |
| 133 | |
| 134 | assert(ind.dtype() == torch::kLong); |
| 135 | assert(size.dtype() == torch::kLong); |
| 136 | |
| 137 | auto yz_size = size.index({"...", Slice(1, 2)}) * size.index({"...", Slice(2, 3)}); // (*, 1) |
| 138 | auto xs = torch::div(ind, yz_size, "floor"); // (*, n) |
| 139 | auto _ind = ind - xs * yz_size; // (*, n) |
| 140 | auto ys = torch::div(_ind, size.index({"...", Slice(2, 3)}), "floor"); // (*, n) |
| 141 | auto zs = _ind - ys * size.index({"...", Slice(2, 3)}); // (*, n) |
| 142 | auto idx = torch::stack({xs, ys, zs}, -1); // (*, n, 3) |
| 143 | return idx; |
| 144 | } |
| 145 | |
| 146 | |
| 147 | // Gather the points belonging to each grid cell. |