torch.narrow() with tensor of start indices
Sophia Terry 
I'm wondering if it's possible to do something in PyTorch...
Basically, what I'm trying to do is something like torch.narrow(), but where the start value is a tensor of indices.
Suppose I had the following tensors:
start_indices = [3, 1, 2, ...] # shape: (M,)
dataset = [ 3, 5, 3, 4, 8, 0, 1, 3, 9, 7, 2, 7, 3, 7, 6, 0, 2, 3, 0, 2, 5, ...
] # shape: (M, N) (here, N = 7, just for demo)In my use case, M ~ 10^3 and N > 10^6
I'm wondering if it's possible to do something like this:
torch.narrow( input=dataset, dim=1, start=start_indices, length=4
)that returns a "narrowed" version of the original dataset, where the start of each row is selected by its corresponding start index. To complete the example I gave above, the result would look like:
result = [ 4, 8, 0, 1, 9, 7, 2, 7, 2, 3, 0, 2, ....
] # shape: (M, 4)I realize that depending on the start indices and length, you could have some lines going out of bounds of the original tensor -> I'm fine with receiving an error in this case.
Right now, I'm currently using a for-loop to do this, but it becomes far too slow as M becomes large. I'm hoping there's a way to do this that takes advantage of data parallelism.
Any ideas? I thought of using torch.index_select(), but it only selects within rows based off of single indices. I want to select within rows by single indices, but then extend those selections out by some fixed length. In that way, what I'm looking for is kind of a "combination" of torch.narrow() and torch.index_select().
Thanks in advance!
Related questions 8 Filling torch tensor with zeros after certain index 0 Pytorch - select region of a tensor using torch function 7 Index a torch tensor with an array Related questions 8 Filling torch tensor with zeros after certain index 0 Pytorch - select region of a tensor using torch function 7 Index a torch tensor with an array 0 How to convert a matrix of torch.tensor to a larger tensor? 3 How can I get argmaxed torch tensor excluding certain index? 2 Efficiently filling torch.Tensor at equal index positions 1 Appling sliding window to torch.tensor and adjusting tensor initial size 4 Split a torch tensor using a same-sized tensor of indices 1 Reduce torch tensor 2 How to use indices from torch.min to index multi-dimensional tensor Load 7 more related questions Show fewer related questions Reset to default
