CODEWITHSUNDEEP
pythonarraysnumpyindexing
user2726995
user2726995

Reputation: 2062

Index Numpy tensor without having to reshape

I have a tensor with the shape (5,48,15). How can I access an element along the 0th axis and still maintain 3 dimensions without needing to reshape. For example:

x.shape                    # this is (5,48,15) 
m = x[0,:,:] 
m.shape                     # This is  (48,15)
m_new = m.reshape(1,48,15) 
m_new.shape                 # This is now (1,48,15)

Is this possible without needing to reshape?

Upvotes: 3

Views: 267

Answers (2)

hpaulj
hpaulj

Reputation: 231335

The selection index needs to be a slice or list (or array):

m = x[[0],:,:]
m = x[:1,:,:]
m = x[0:1,:,:]

Upvotes: 1

Alex Riley
Alex Riley

Reputation: 176740

When you index an axis with a single integer, as with x[0, :, :], the dimensionality of the returned array drops by one.

To keep three dimensions, you can either...

  • insert a new axis at the same time as indexing:

    >>> x[None, 0, :, :].shape
(1, 48, 15)

  • or use slicing:

    >>> x[:1, :, :].shape
(1, 48, 15)

  • or use fancy indexing:

    >>> x[[0], :, :].shape
(1, 48, 15)

Upvotes: 4

Related Questions