code: def dot_or_identity(A, B, device=None):
      # if A is None, treat as identity matrix
      if A is None:
          return B
      elif len(A.shape) == 1:
          if device is None:
              return B[A]
          else:
              return B[A].to(device)
      else:
          return A @ B
  
imports: import dgl.function
  import dgl.nn.pytorch
  import torch
  import torch.nn
  from torch.nn import init
  from utils import get_activation, to_etype_name
label_desc: Use image node_id 1 for calling a global function with example usage: dot_or_identity(A, B, device) and returns: B, unknown, B, B