[Feature Request] batchwise gather_nd/scatter_nd

[ShownX]

Description

Assuming I have a data tensor X and indices tensor Y: X has a shape (batch, N, 3) and Y has a shape (batch, 1, M), I hope there is a new operation can do batch-wise gather_nd and scatter_nd using the index within the batch and gather the data to new tensor (batch, M, 3)

Implementation

Forwarding...

gathered = []
for i in range(X.shape[0]):
     gathered.append(nd.gather_nd(X[i], Y[i]))
return nd.concat(*gathered)
``

[sxjscience]

This should be the correct version of batch_take. We need to support this in the future.

[Cogito2012]

@ShownX Have you found the solution to this problem? I am also struggling on this issue...