Could you provide a more general version to fulfill tf.gather_nd, not only limited to 2d tensor?
For example, if I want to gather 3d tensor with size CxHxW by giving the indices with size Nx2, where the indices values correspond to the coordinates in HxW grid, how can I get the results with size CxN?
Actually in tensorflow, this can be easily implemented by tf.gather_nd…