code: def bmm_maybe_select(A, B, index):
      """Slice submatrices of A by the given index and perform bmm.
  
      B is a 3D tensor of shape (N, D1, D2), which can be viewed as a stack of
      N matrices of shape (D1, D2). The input index is an integer vector of length M.
      A could be either:
      (1) a dense tensor of shape (M, D1),
      (2) an integer vector of length M.
      The result C is a 2D matrix of shape (M, D2)
  
      For case (1), C is computed by bmm:
      ::
  
          C[i, :] = matmul(A[i, :], B[index[i], :, :])
  
      For case (2), C is computed by index select:
      ::
  
          C[i, :] = B[index[i], A[i], :]
  
      Parameters
      ----------
      A : mxnet.NDArray
          lhs tensor
      B : mxnet.NDArray
          rhs tensor
      index : mxnet.NDArray
          index tensor
  
      Returns
      -------
      C : mxnet.NDArray
          return tensor
      """
      if A.dtype in (np.int32, np.int64) and len(A.shape) == 1:
          return B[index, A, :]
      else:
          BB = nd.take(B, index, axis=0)
          return nd.batch_dot(A.expand_dims(1), BB).squeeze(1)
  
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 2 for calling a global function with example usage: bmm_maybe_select(A, B, index) and returns: B, nd