code: def dot_or_identity(A, B, device=None):
      """
      Return as identity matrix if A is none.
      If A exists, return as dot product.
      """
      if A is None:
          return B
      if len(A.shape) == 1:
          if device is None:
              return B[A]
          return B[A].to(device)
      return A @ B
  
imports: import dgl.function
  import dgl.nn.pytorch
  from dgl import DGLError
  import torch
  import torch.nn
  from torch.nn import init
  from .utils import get_activation
label_desc: Use image node_id 1 for calling a global function with example usage: dot_or_identity(A, B, device) and returns: unknown, B, B, B