I would like to flatten a tensor into column-major order. What is the optimal way for achieving this in PyTorch? Row-major order seems to be the default in PyTorch’s flatten function, and I don’t think there is an order option like in Numpy’s flatten function. Thanks!
Use tranpose method to reshape your tensor then flatten i.e.
x = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
x.flatten()
tensor([1, 2, 3, 4, 5, 6, 7, 8, 9])
x.transpose(1, 0).flatten()
tensor([1, 4, 7, 2, 5, 8, 3, 6, 9])
Thank you for your reply!
That works well for matrices (where the transpose can be written as simply x.t().flatten()
)
However, I am dealing with ndarrays, so that a matrix transpose is not applicable.
I just found that the solution had already been posted on this forum, here:
The answer was:
t = torch.rand(3, 3, 3)
# convert to column-major order
t.set_(t.storage(), t.storage_offset(), t.size(), tuple(reversed(t.stride())))
t.flatten() # 1D array in column-major order
Note that if you just want a tensor’s 1D representation in column-major order, the above operation will change the ordering of the elements in tensor t
. This function will pull out just the flattened array in column-major order:
def flatten_fortran_w_clone(t):
nt = t.clone()
nt = nt.set_(nt.storage(), nt.storage_offset(), nt.size(), tuple(reversed(nt.stride())))
return nt.flatten()
But there might be a better way of doing it without having to create a copy…
Hi @brudfors,
Your flatten_fortran_w_clone
functiton doesn’t work for a 2D-array, for example,
def flatten_fortran_w_clone(t):
nt = t.clone()
nt = nt.set_(nt.storage(), nt.storage_offset(), nt.size(), tuple(reversed(nt.stride())))
return nt.flatten()
x=torch.arange(6).reshape(3,2)
#returns tensor([[0, 1],
# [2, 3],
# [4, 5]])
flatten_fortran_w_clone(x) #returns tensor([0, 2, 1, 3, 2, 4]) #incorrect
x.transpose(-2,-1).flatten() #returns tensor([0, 2, 4, 1, 3, 5]) #correct
Also, is there a way to invert this process? For example, column-major flatten an array then reshape (via column-major formatting, instead of row-major formatting) to get back to the original array?
Also, for anyone reading this nt.storage()
is now deprecated and should be replaced with nt.untyped_storage()
instead, i.e.,
def flatten_fortran_w_clone(t):
nt = t.clone()
nt = nt.set_(nt.untyped_storage(), nt.storage_offset(), nt.size(), tuple(reversed(nt.stride())))
return nt.flatten()