Class VarLenTensor

pixal3d/modules/sparse/basic.py:17–302

Sequential tensor with variable length. Args: feats (torch.Tensor): Features of the varlen tensor. layout (List[slice]): Layout of the varlen tensor for each batch

class VarLenTensor:
    """
    Sequential tensor with variable length.
    
    Args:
        feats (torch.Tensor): Features of the varlen tensor.
        layout (List[slice]): Layout of the varlen tensor for each batch
    """
    def __init__(self, feats: torch.Tensor, layout: List[slice]=None):
        self.feats = feats
        self.layout = layout if layout is not None else [slice(0, feats.shape[0])]
        self._cache = {}
        
    @staticmethod
    def layout_from_seqlen(seqlen: list) -> List[slice]:
        """
        Create a layout from a tensor of sequence lengths.
        """
        layout = []
        start = 0
        for l in seqlen:
            layout.append(slice(start, start + l))
            start += l
        return layout
        
    @staticmethod
    def from_tensor_list(tensor_list: List[torch.Tensor]) -> 'VarLenTensor':
        """
        Create a VarLenTensor from a list of tensors.
        """
        feats = torch.cat(tensor_list, dim=0)
        layout = []
        start = 0
        for tensor in tensor_list:
            layout.append(slice(start, start + tensor.shape[0]))
            start += tensor.shape[0]
        return VarLenTensor(feats, layout)
    
    def to_tensor_list(self) -> List[torch.Tensor]:
        """
        Convert a VarLenTensor to a list of tensors.
        """
        tensor_list = []
        for s in self.layout:
            tensor_list.append(self.feats[s])
        return tensor_list
    
    def __len__(self) -> int:
        return len(self.layout)
    
    @property
    def shape(self) -> torch.Size:
        return torch.Size([len(self.layout), *self.feats.shape[1:]])
    
    def dim(self) -> int:
        return len(self.shape)
    
    @property