code: def _gather_nd(index, src):
      """Similar to PyTorch's gather nd on first dimension."""
      ctx = context(src)
      shp = index.shape
      ndim = src.ndim
      offsets = []
      stride = 1
      for i in reversed(range(1, ndim)):
          di = shp[i]
          offset_i = nd.arange(di, dtype=index.dtype)
          offsets.append(
              (stride * offset_i).reshape(
                  (1,) * i + (di,) + (1,) * (ndim - 1 - i)
              )
          )
          stride *= di
      if ndim > 1:
          new_idx = index * stride + copy_to(sum(offsets), ctx)
      else:
          new_idx = index
      src = src.reshape(-1)
      new_idx = new_idx.reshape(-1)
      rst = nd.take(src, new_idx).reshape(shp)
      return rst
  
imports: import mxnet
  import numpy
  from mxnet import nd
  from ..._sparse_ops import _bwd_segment_cmp, _csrmask, _csrmm, _csrsum, _gsddmm, _gspmm, _scatter_add, _segment_reduce
  from ...base import ALL, dgl_warning, is_all
  from ...heterograph_index import create_unitgraph_from_csr
  from .tensor import asnumpy, context, copy_to, to_backend_ctx, zerocopy_from_numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _gather_nd(index, src) and returns: rst