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 : torch.Tensor
          lhs tensor
      B : torch.Tensor
          rhs tensor
      index : torch.Tensor
          index tensor
  
      Returns
      -------
      C : torch.Tensor
          return tensor
      """
      if A.dtype == th.int64 and len(A.shape) == 1:
          # following is a faster version of B[index, A, :]
          B = B.view(-1, B.shape[2])
          flatidx = index * B.shape[1] + A
          return B.index_select(0, flatidx)
      else:
          BB = B.index_select(0, index)
          return th.bmm(A.unsqueeze(1), BB).squeeze()
  
imports: import torch
  import torch.nn.functional
  from torch import nn
  from ...None import DGLGraph, function
  from ...base import dgl_warning
label_desc: Use image node_id 2 for calling a global function with example usage: bmm_maybe_select(A, B, index) and returns: B, th