Given a dict or a list containing tensor, reshape dimension `start` to `end` to the new shape
(
arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]],
start: int = 0,
end: int = -1, # included
shape: T.Union[int, T.List[int], T.Tuple[int]] = None,
)
| 159 | |
| 160 | |
| 161 | def reshape( |
| 162 | arr: T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]], |
| 163 | start: int = 0, |
| 164 | end: int = -1, # included |
| 165 | shape: T.Union[int, T.List[int], T.Tuple[int]] = None, |
| 166 | ) -> T.Union[torch.Tensor, T.List[torch.Tensor], T.Dict[str, T.Any]]: |
| 167 | """ |
| 168 | Given a dict or a list containing tensor, |
| 169 | reshape dimension `start` to `end` to the new shape |
| 170 | """ |
| 171 | shape = list(shape) |
| 172 | |
| 173 | if isinstance(arr, torch.Tensor): |
| 174 | arr_shape = arr.shape |
| 175 | if end < 0: |
| 176 | end = end + arr.ndim |
| 177 | new_shape = list(arr_shape[:start]) + list(shape) + list(arr_shape[end + 1:]) |
| 178 | arr = arr.reshape(*new_shape) |
| 179 | return arr |
| 180 | elif isinstance(arr, (list, tuple)): |
| 181 | return [reshape(x, start=start, end=end, shape=shape) for x in arr] |
| 182 | elif isinstance(arr, dict): |
| 183 | for key, val in arr.items(): |
| 184 | arr[key] = reshape(val, start=start, end=end, shape=shape) |
| 185 | return arr |
| 186 | else: |
| 187 | return arr |
| 188 | |
| 189 | |
| 190 | def create_pcd( |