Reverse operation of torch.cat

Viewed 1525

Suppose I have a tensor like [A,A,A,A...,A].

How can I quickly obtain [[A],[A],[A],...,[A]] as a tensor in torch?

3 Answers

You can use torch.chunk as the inverse of cat, but it looks like you want unsqueeze(1):

A = torch.randn(2, 3)
A_rep = (A, A, A, A, A, A, A, A)
catted = torch.cat(A_rep)

#uncatted = torch.chunk(catted, len(A_rep))
catted.unsqueeze(1)

From the doc, use split or reshape.

>>> a = torch.arange(10).reshape(5,2)
>>> a
tensor([[0, 1],
        [2, 3],
        [4, 5],
        [6, 7],
        [8, 9]])
>>> torch.split(a, 2)
(tensor([[0, 1],
         [2, 3]]),
 tensor([[4, 5],
         [6, 7]]),
 tensor([[8, 9]]))
>>> torch.split(a, [1,4])
(tensor([[0, 1]]),
 tensor([[2, 3],
         [4, 5],
         [6, 7],
         [8, 9]]))

You are looking for torch.unbind which does exactly what you want.

import torch

tensor = torch.rand(3, 4, 5)  # tensor of shape (3, 4, 5)
l = tensor.unbind(dim=0)  # list of 3 tensors of shape (4, 5)

You can unbind along whatever dimension you want, default is dim 0. You can use python destructuring to quickly get the tensors in variables:

a, b, c = tensor.unbind(dim=0)  # 3 tensors of shape (4, 5)
Related