Similar operation like `numpy.take`

1

Is there any function same as the numpy.take function ?

a = np.random.rand(10,2)
a
array([[0.14009622, 0.51205327],
[0.79369912, 0.72729038],
[0.7041849 , 0.38336232],
[0.24226114, 0.82128616],
[0.79437292, 0.80704659],
[0.49525265, 0.94109091],
[0.68034488, 0.17979233],
[0.49363128, 0.1746095 ],
[0.01659873, 0.0146231 ],
[0.00280727, 0.04646634]])
#if we want to take the [1,2,3] rows
np.take(a,[1,2,3],axis=0)
array([[0.79369912, 0.72729038],
[0.7041849 , 0.38336232],
[0.24226114, 0.82128616]])

2

Do you need something like - Select rows from a 2D tensor

3

na=torch.randn(1, 1, 3, 3)
print(rna)
torch.take(rna,torch.tensor([4]))