class_docstr: ## Long Short-Term Memory Cell
  
  LSTM Cell computes $c$, and $h$. $c$ is like the long-term memory,
  and $h$ is like the short term memory.
  We use the input $x$ and $h$ to update the long term memory.
  In the update, some features of $c$ are cleared with a forget gate $f$,
  and some features $i$ are added through a gate $g$.
  
  The new short term memory is the $  anh$ of the long-term memory
  multiplied by the output gate $o$.
  
  Note that the cell doesn't look at long term memory $c$ when doing the update. It only modifies it.
  Also $c$ never goes through a linear transformation.
  This is what solves vanishing and exploding gradients.
  
  Here's the update rule.
  
  egin{align}
  c_t &= \sigma(f_t) \odot c_{t-1} + \sigma(i_t) \odot        anh(g_t) \
  h_t &= \sigma(o_t) \odot    anh(c_t)
  \end{align}
  
  $\odot$ stands for element-wise multiplication.
  
  Intermediate values and gates are computed as linear transformations of the hidden
  state and input.
  
  egin{align}
  i_t &= lin_x^i(x_t) + lin_h^i(h_{t-1}) \
  f_t &= lin_x^f(x_t) + lin_h^f(h_{t-1}) \
  g_t &= lin_x^g(x_t) + lin_h^g(h_{t-1}) \
  o_t &= lin_x^o(x_t) + lin_h^o(h_{t-1})
  \end{align}
class_docstr_tok: ["#","#","Long","Short-Term","Memory","Cell","LSTM","Cell","computes","$","c","$",",","and","$","h","$",".","$","c","$","is","like","the","long-term","memory",",","and","$","h","$","is","like","the","short","term","memory",".","We","use","the","input","$","x","$","and","$","h","$","to","update","the","long","term","memory",".","In","the","update",",","some","features","of","$","c","$","are","cleared","with","a","forget","gate","$","f","$",",","and","some","features","$","i","$","are","added","through","a","gate","$","g","$",".","The","new","short","term","memory","is","the","$","anh","$","of","the","long-term","memory","multiplied","by","the","output","gate","$","o","$",".","Note","that","the","cell","does","n't","look","at","long","term","memory","$","c","$","when","doing","the","update",".","It","only","modifies","it",".","Also","$","c","$","never","goes","through","a","linear","transformation",".","This","is","what","solves","vanishing","and","exploding","gradients",".","Here","'s","the","update","rule",".","\begin","{","align","}","c_t","&","=","\\sigma","(","f_t",")","\\odot","c_","{","t-1","}","+","\\sigma","(","i_t",")","\\odot","anh","(","g_t",")","\\","h_t","&","=","\\sigma","(","o_t",")","\\odot","anh","(","c_t",")","\\end","{","align","}","$","\\odot","$","stands","for","element-wise","multiplication",".","Intermediate","values","and","gates","are","computed","as","linear","transformations","of","the","hidden","state","and","input",".","\begin","{","align","}","i_t","&","=","lin_x^i","(","x_t",")","+","lin_h^i","(","h_","{","t-1","}",")","\\","f_t","&","=","lin_x^f","(","x_t",")","+","lin_h^f","(","h_","{","t-1","}",")","\\","g_t","&","=","lin_x^g","(","x_t",")","+","lin_h^g","(","h_","{","t-1","}",")","\\","o_t","&","=","lin_x^o","(","x_t",")","+","lin_h^o","(","h_","{","t-1","}",")","\\end","{","align","}"]
code: def __init__(
      self, input_size: int, hidden_size: int, layer_norm: bool = False
  ):
      super().__init__()
  
      # These are the linear layer to transform the `input` and `hidden` vectors.
      # One of them doesn't need a bias since we add the transformations.
  
      # This combines $lin_x^i$, $lin_x^f$, $lin_x^g$, and $lin_x^o$ transformations.
      self.hidden_lin = nn.Linear(hidden_size, 4 * hidden_size)
      # This combines $lin_h^i$, $lin_h^f$, $lin_h^g$, and $lin_h^o$ transformations.
      self.input_lin = nn.Linear(
          input_size, 4 * hidden_size, bias=False
      )
  
      # Whether to apply layer normalizations.
      #
      # Applying layer normalization gives better results.
      # $i$, $f$, $g$ and $o$ embeddings are normalized and $c_t$ is normalized in
      # $h_t = o_t \odot \tanh(\mathop{LN}(c_t))$
      if layer_norm:
          self.layer_norm = nn.ModuleList(
              [nn.LayerNorm(hidden_size) for _ in range(4)]
          )
          self.layer_norm_c = nn.LayerNorm(hidden_size)
      else:
          self.layer_norm = nn.ModuleList(
              [nn.Identity() for _ in range(4)]
          )
          self.layer_norm_c = nn.Identity()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/lstm/__init__.py
imports: from typing import Optional, Tuple
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 1 to create a new LSTMCell object from inherited base classes: Module with example: obj = LSTMCell(input_size, hidden_size, layer_norm)