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.
docstr: Updates the memory with newly encountered interactions
  :obj:`(src, dst, t, raw_msg)`.
docstr_tok: ["Updates","the","memory","with","newly","encountered","interactions",":","obj",":","`","(","src",",","dst",",","t",",","raw_msg",")","`","."]
code: def update_state(
      self, src: Tensor, dst: Tensor, t: Tensor, raw_msg: Tensor
  ):
      """Updates the memory with newly encountered interactions
      :obj:`(src, dst, t, raw_msg)`.
      """
      n_id = torch.cat([src, dst]).unique()
  
      if self.training:
          self._update_memory(n_id)
          self._update_msg_store(src, dst, t, raw_msg, self.msg_s_store)
          self._update_msg_store(dst, src, t, raw_msg, self.msg_d_store)
      else:
          self._update_msg_store(src, dst, t, raw_msg, self.msg_s_store)
          self._update_msg_store(dst, src, t, raw_msg, self.msg_d_store)
          self._update_memory(n_id)
  
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 7 for calling the TGNMemory obj's underlying member method code with example usage: obj.update_state(src, dst, t, raw_msg) without return types