code: def matmul_maybe_select(A, B):
      """Perform Matrix multiplication C = A * B but A could be an integer id vector.
  
      If A is an integer vector, we treat it as multiplying a one-hot encoded tensor.
      In this case, the expensive dense matrix multiply can be replaced by a much
      cheaper index lookup.
  
      For example,
      ::
  
          A = [2, 0, 1],
          B = [[0.1, 0.2],
               [0.3, 0.4],
               [0.5, 0.6]]
  
      then matmul_maybe_select(A, B) is equivalent to
      ::
  
          [[0, 0, 1],     [[0.1, 0.2],
           [1, 0, 0],  *   [0.3, 0.4],
           [0, 1, 0]]      [0.5, 0.6]]
  
      In all other cases, perform a normal matmul.
  
      Parameters
      ----------
      A : mxnet.NDArray
          lhs tensor
      B : mxnet.NDArray
          rhs tensor
  
      Returns
      -------
      C : mxnet.NDArray
          result tensor
      """
      if A.dtype in (np.int32, np.int64) and len(A.shape) == 1:
          return nd.take(B, A, axis=0)
      else:
          return nd.dot(A, B)
  
file_path: dgl/python/dgl/nn/mxnet/utils.py
imports: import numpy
  from mxnet import gluon, nd
  from ...None import DGLGraph
label_desc: Use image node_id 1 for calling a global function with example usage: matmul_maybe_select(A, B) and returns: nd, nd