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 : torch.Tensor
          lhs tensor
      B : torch.Tensor
          rhs tensor
  
      Returns
      -------
      C : torch.Tensor
          result tensor
      """
      if A.dtype == th.int64 and len(A.shape) == 1:
          return B.index_select(0, A)
      else:
          return th.matmul(A, B)
  
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 1 for calling a global function with example usage: matmul_maybe_select(A, B) and returns: B, th