Indexing whole Tensor along specific dimension and specific channels
Let say we have a Tensor A with the dimension dim(A)=[i, j, k=6, u, v]. Now we are interested to get the whole tensor at dimension k with channels=[0:3]. I know we can get it this way:
B = A[:, :, 0:3, :, :]

Now I would like to know if there is any better &quot;pythonic&quot; way to achieve the same result without doing this suboptimal indexing. I mean something like.
B = subset(A, dim=2, index=[0, 1, 2])

No matter in which framework, i.e. pytorch, tensorflow, numpy, etc.
Thanks a lot

In numpy, you can use the take method:
B = A.take([0,1,2], axis=2)


In TensorFlow, there is not really a more concise way than using the traditionnal approach. Using tf.slice would be really verbose:
B = tf.slice(A,[0,0,0,0,0],[-1,-1,3,-1,-1])

You can potentially use the experimental version of take (since TF 2.4):
B = tf.experimental.numpy.take(A, [0,1,2], axis=2)


in PyTorch, you can use index_select:
torch.index_select(A, dim=2, index=torch.tensor([0,1,2]))


Note that you can skip listing explicitly the first dimensions (or the last) by using an ellipsis:
# Both are equivalent in that case
B = A[..., 0:3, :, :]
B = A[:, :, 0:3, ...]

