Given x y z index, change to the linear index. (matlab's sub2ind) Args: idx: (*, n, 3) int64_t size: (*, 3) int64_t Returns: (*, n) int64_t
| 99 | // (*, n) int64_t |
| 100 | // |
| 101 | torch::Tensor sub2ind( |
| 102 | const torch::Tensor & idx, // (*, n, 3) int64_t |
| 103 | const torch::Tensor & size // (*, 3) int64_t |
| 104 | ) { |
| 105 | |
| 106 | assert(idx.dtype() == torch::kLong); |
| 107 | assert(size.dtype() == torch::kLong); |
| 108 | |
| 109 | auto linear_idx = |
| 110 | idx.index({"...", 2}) + |
| 111 | idx.index({"...", 1}) * size.index({"...", Slice(2, 3)}) + |
| 112 | idx.index({"...", 0}) * (size.index({"...", Slice(1, 2)}) * size.index({"...", Slice(2, 3)})); // (*, n) |
| 113 | return linear_idx; |
| 114 | } |
| 115 | |
| 116 | |
| 117 |