(self, idx)
| 236 | return self.__elemwise__(other, lambda x, y: torch.div(y, x)) |
| 237 | |
| 238 | def __getitem__(self, idx): |
| 239 | if isinstance(idx, int): |
| 240 | idx = [idx] |
| 241 | elif isinstance(idx, slice): |
| 242 | idx = range(*idx.indices(self.shape[0])) |
| 243 | elif isinstance(idx, list): |
| 244 | assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}" |
| 245 | elif isinstance(idx, torch.Tensor): |
| 246 | if idx.dtype == torch.bool: |
| 247 | assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}" |
| 248 | idx = idx.nonzero().squeeze(1) |
| 249 | elif idx.dtype in [torch.int32, torch.int64]: |
| 250 | assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}" |
| 251 | else: |
| 252 | raise ValueError(f"Unknown index type: {idx.dtype}") |
| 253 | else: |
| 254 | raise ValueError(f"Unknown index type: {type(idx)}") |
| 255 | |
| 256 | new_feats = [] |
| 257 | new_layout = [] |
| 258 | start = 0 |
| 259 | for new_idx, old_idx in enumerate(idx): |
| 260 | new_feats.append(self.feats[self.layout[old_idx]]) |
| 261 | new_layout.append(slice(start, start + len(new_feats[-1]))) |
| 262 | start += len(new_feats[-1]) |
| 263 | new_feats = torch.cat(new_feats, dim=0).contiguous() |
| 264 | new_tensor = VarLenTensor(feats=new_feats, layout=new_layout) |
| 265 | return new_tensor |
| 266 | |
| 267 | def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| 268 | if isinstance(dim, int): |
nothing calls this directly
no test coverage detected