Tensorrt: Q: is there a layer that can do "select by index ?"

[taomiao]

Description

assume :
X = [ 2, 3, 4, 5, 6, 7, 8]
idx = [ 2, 3 ]
S_X = X[idx] ===> S_X = [4, 5]

Environment

TensorRT Version:
GPU Type:
Nvidia Driver Version:
CUDA Version:
CUDNN Version:
Operating System + Version:
Python Version (if applicable):
TensorFlow Version (if applicable):
PyTorch Version (if applicable):
Baremetal or Container (if container which image + tag):

Relevant Files

Steps To Reproduce

[1-800-BAD-CODE]

I'll second that it would be useful to have in a future release an IIndexSelectLayer that mimics torch.index_select. It would provide a key part of some beam search algorithms.

[rmccorm4]

@pranavm-nvidia Can this be achieved with a Select layer or some other workaround?

[1-800-BAD-CODE]

Although the documentation is limited, it looks like an IGatherLayer is essentially an index_select. I'll try it this week.

[rmccorm4]

it looks like an IGatherLayer is essentially an index_select. I'll try it this week.

In case it helps: https://github.com/NVIDIA/TensorRT/blob/master/samples/opensource/sampleNMT/model/softmaxLikelihood.cpp#L75-L79

[pranavm-nvidia]

Yeah, Gather is what you want here

[rmccorm4]

@1-800-BAD-CODE did it work out for you?

[1-800-BAD-CODE]

@1-800-BAD-CODE did it work out for you?

Yes it seems to be the correct solution here.