Unbind a sparse tensor along a dimension. Args: input (SparseTensor): Sparse tensor to unbind. dim (int): Dimension to unbind.
(input: SparseTensor, dim: int)
| 822 | |
| 823 | |
| 824 | def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]: |
| 825 | """ |
| 826 | Unbind a sparse tensor along a dimension. |
| 827 | |
| 828 | Args: |
| 829 | input (SparseTensor): Sparse tensor to unbind. |
| 830 | dim (int): Dimension to unbind. |
| 831 | """ |
| 832 | if dim == 0: |
| 833 | return [input[i] for i in range(input.shape[0])] |
| 834 | else: |
| 835 | feats = input.feats.unbind(dim) |
| 836 | return [input.replace(f) for f in feats] |