MCPcopy Index your code
hub / github.com/microsoft/TRELLIS / __getitem__

Method __getitem__

trellis/modules/sparse/basic.py:344–368  ·  view source on GitHub ↗
(self, idx)

Source from the content-addressed store, hash-verified

342 return self.__elemwise__(other, lambda x, y: torch.div(y, x))
343
344 def __getitem__(self, idx):
345 if isinstance(idx, int):
346 idx = [idx]
347 elif isinstance(idx, slice):
348 idx = range(*idx.indices(self.shape[0]))
349 elif isinstance(idx, torch.Tensor):
350 if idx.dtype == torch.bool:
351 assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
352 idx = idx.nonzero().squeeze(1)
353 elif idx.dtype in [torch.int32, torch.int64]:
354 assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
355 else:
356 raise ValueError(f"Unknown index type: {idx.dtype}")
357 else:
358 raise ValueError(f"Unknown index type: {type(idx)}")
359
360 coords = []
361 feats = []
362 for new_idx, old_idx in enumerate(idx):
363 coords.append(self.coords[self.layout[old_idx]].clone())
364 coords[-1][:, 0] = new_idx
365 feats.append(self.feats[self.layout[old_idx]])
366 coords = torch.cat(coords, dim=0).contiguous()
367 feats = torch.cat(feats, dim=0).contiguous()
368 return SparseTensor(feats=feats, coords=coords)
369
370 def register_spatial_cache(self, key, value) -> None:
371 """

Callers

nothing calls this directly

Calls 1

SparseTensorClass · 0.85

Tested by

no test coverage detected