code: def select_vectors_by_position(data, positions):
      """Select each batch with the given positions.
  
      Once advanced indexing can be hybridized, we can revise the implementation.
  
      out[i, j, ...] = data[i, positions[i, j], ...]
  
      Parameters
      ----------
      data
          Input tensor of contextualized token embeddings
          Shape (batch_size, seq_length, ...)
      positions
          Input tensor of the positions.
          Shape (batch_size, num_sel_positions).
          For each sample in the batch, the values in this tensor must not exceed
          the length of the sequence.
  
      Returns
      -------
      out
          The selection result.
          Shape (batch_size, num_sel_positions, ...)
      """
      # Here, we use gather_nd to select the output from data:
      # Need to compute
      #   out[i, j, :] = in[i, masked_position[i, j], :]
      # Thus, construct a indices with shape [2, batch_size, num_masked_position], where
      #     indices[0, i, j] = i
      #     indices[1, i, j] = masked_position[i, j]
      # Then, out = gather_nd(in, indices)
      positions = positions.astype(np.int32)
      # batch_idx.shape = (batch_size, 1) as [[0], [1], [2], ...]
      batch_idx = np.expand_dims(
          npx.arange_like(positions, axis=0), axis=1
      ).astype(np.int32)
      batch_idx = batch_idx + np.zeros_like(positions)
      indices = np.stack([batch_idx, positions])
      # TODO(sxjscience) We can revise the implementation to advanced indexing
      #  once the bug in MXNet is solved:
      #  https://github.com/apache/incubator-mxnet/issues/18919
      out = npx.gather_nd(data, indices)
      return out
  
file_path: gluon-nlp/src/gluonnlp/op.py
imports: import mxnet
  from mxnet import np, npx
  import math
  from mxnet import use_np
label_desc: Use image node_id 1 for calling a global function with example usage: select_vectors_by_position(data, positions) and returns: out