Obtaining & Assigning values with tensors based on multiple indexes

1

#1 Problems come up when obtaining values with two index vector from a tensor.
For example:

>>> import torch
>>> torch.__version__
'2.4.0'
>>> x=torch.randn(2,3,4,5)
>>> a=[0]
>>> b=[0,1]
>>> x[a,:,b,:].shape
torch.Size([2, 3, 5])
>>> x[a][:,:,b,:].shape
torch.Size([1, 3, 2, 5])

Two index at different dimension x[a,:,b,:] doesn’t return the expected shape of tensor as x[a][:,:,b,:]. Furthermore:

>>> a=[0,1]
>>> b=[0,1,2]
>>> x[a,:,b,:]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [2], [3]

The mechanism of multiple index quite confuses me.

#2 Even though I can obtain the expected shape of tensor from x[a][:,:,b,:], however, assigning values to it is no use. Specifically:

>>> a=[0]
>>> b=[0,1]
>>> y=torch.zeros(1,3,2,5)
>>> x[a][:,:,b,:]=y
>>> x[a][:,:,b,:]
tensor([[[[-0.4322, -0.2585, -1.2516, -0.6877, -0.2141],
          [ 1.9328, -1.1326,  2.1536, -0.1765,  0.1218]],

         [[-0.0773,  1.2094, -0.5099,  0.7909, -0.1155],
          [-0.1638,  1.0042,  0.2794,  0.2649,  1.0520]],

         [[ 0.5747,  0.7239,  0.9941, -1.0764, -0.3551],
          [-0.3900,  0.7646, -0.1625,  0.5727, -1.6328]]]])

The values in tensor x doesn’t become y. I guess it’s because the x[a][:,:,b,:] return a new view instead of reference.
So how could I assign y to x correctly? :melting_face: