class_docstr: The Temporal Graph Network (TGN) memory model from the
  `"Temporal Graph Networks for Deep Learning on Dynamic Graphs"
  <https://arxiv.org/abs/2006.10637>`_ paper.
  
  .. note::
  
      For an example of using TGN, see `examples/tgn.py
      <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
      tgn.py>`_.
  
  Args:
      num_nodes (int): The number of nodes to save memories for.
      raw_msg_dim (int): The raw message dimensionality.
      memory_dim (int): The hidden memory dimensionality.
      time_dim (int): The time encoding dimensionality.
      message_module (torch.nn.Module): The message function which
          combines source and destination node memory embeddings, the raw
          message and the time encoding.
      aggregator_module (torch.nn.Module): The message aggregator function
          which aggregates messages to the same destination into a single
          representation.
code: def _compute_msg(
      self,
      n_id: Tensor,
      msg_store: TGNMessageStoreType,
      msg_module: Callable,
  ):
      data = [msg_store[i] for i in n_id.tolist()]
      src, dst, t, raw_msg = list(zip(*data))
      src = torch.cat(src, dim=0)
      dst = torch.cat(dst, dim=0)
      t = torch.cat(t, dim=0)
      raw_msg = torch.cat(raw_msg, dim=0)
      t_rel = t - self.last_update[src]
      t_enc = self.time_enc(t_rel.to(raw_msg.dtype))
  
      msg = msg_module(
          self.memory[src], self.memory[dst], raw_msg, t_enc
      )
  
      return msg, t, src, dst
  
imports: import copy
  from typing import Callable, Dict, Tuple
  import torch
  from torch import Tensor
  from torch.nn import GRUCell, Linear
  from torch_geometric.nn.inits import zeros
  from torch_geometric.utils import scatter
  from torch_geometric.utils._scatter import scatter_argmax
label_desc: Use image node_id 12 for calling the TGNMemory obj's underlying member method code with example usage: obj._compute_msg(n_id, msg_store, msg_module) and returns: msg,  t,  src,  dst