(self, idx)
| 730 | return new_tensor |
| 731 | |
| 732 | def __getitem__(self, idx): |
| 733 | if isinstance(idx, int): |
| 734 | idx = [idx] |
| 735 | elif isinstance(idx, slice): |
| 736 | idx = range(*idx.indices(self.shape[0])) |
| 737 | elif isinstance(idx, list): |
| 738 | assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}" |
| 739 | elif isinstance(idx, torch.Tensor): |
| 740 | if idx.dtype == torch.bool: |
| 741 | assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}" |
| 742 | idx = idx.nonzero().squeeze(1) |
| 743 | elif idx.dtype in [torch.int32, torch.int64]: |
| 744 | assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}" |
| 745 | else: |
| 746 | raise ValueError(f"Unknown index type: {idx.dtype}") |
| 747 | else: |
| 748 | raise ValueError(f"Unknown index type: {type(idx)}") |
| 749 | |
| 750 | new_coords = [] |
| 751 | new_feats = [] |
| 752 | new_layout = [] |
| 753 | new_shape = torch.Size([len(idx)] + list(self.shape[1:])) |
| 754 | start = 0 |
| 755 | for new_idx, old_idx in enumerate(idx): |
| 756 | new_coords.append(self.coords[self.layout[old_idx]].clone()) |
| 757 | new_coords[-1][:, 0] = new_idx |
| 758 | new_feats.append(self.feats[self.layout[old_idx]]) |
| 759 | new_layout.append(slice(start, start + len(new_coords[-1]))) |
| 760 | start += len(new_coords[-1]) |
| 761 | new_coords = torch.cat(new_coords, dim=0).contiguous() |
| 762 | new_feats = torch.cat(new_feats, dim=0).contiguous() |
| 763 | new_tensor = SparseTensor(feats=new_feats, coords=new_coords, shape=new_shape) |
| 764 | new_tensor.register_spatial_cache('layout', new_layout) |
| 765 | return new_tensor |
| 766 | |
| 767 | def clear_spatial_cache(self) -> None: |
| 768 | """ |
nothing calls this directly
no test coverage detected