Unbind a sparse tensor along a dimension. Args: input (SparseTensor): Sparse tensor to unbind. dim (int): Dimension to unbind.
(input: SparseTensor, dim: int)
| 445 | |
| 446 | |
| 447 | def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]: |
| 448 | """ |
| 449 | Unbind a sparse tensor along a dimension. |
| 450 | |
| 451 | Args: |
| 452 | input (SparseTensor): Sparse tensor to unbind. |
| 453 | dim (int): Dimension to unbind. |
| 454 | """ |
| 455 | if dim == 0: |
| 456 | return [input[i] for i in range(input.shape[0])] |
| 457 | else: |
| 458 | feats = input.feats.unbind(dim) |
| 459 | return [input.replace(f) for f in feats] |