class_docstr: ## HyperLSTM Cell
  
  For HyperLSTM the smaller network and the larger network both have the LSTM structure.
  This is defined in Appendix A.2.2 in the paper.
class_docstr_tok: ["#","#","HyperLSTM","Cell","For","HyperLSTM","the","smaller","network","and","the","larger","network","both","have","the","LSTM","structure",".","This","is","defined","in","Appendix","A.2.2","in","the","paper","."]
docstr: `input_size` is the size of the input $x_t$,
  `hidden_size` is the size of the LSTM, and
  `hyper_size` is the size of the smaller LSTM that alters the weights of the larger outer LSTM.
  `n_z` is the size of the feature vectors used to alter the LSTM weights.
  
  We use the output of the smaller LSTM to compute $z_h^{i,f,g,o}$, $z_x^{i,f,g,o}$ and
  $z_b^{i,f,g,o}$ using linear transformations.
  We calculate $d_h^{i,f,g,o}(z_h^{i,f,g,o})$, $d_x^{i,f,g,o}(z_x^{i,f,g,o})$, and
  $d_b^{i,f,g,o}(z_b^{i,f,g,o})$ from these, using linear transformations again.
  These are then used to scale the rows of weight and bias tensors of the main LSTM.
  
  📝 Since the computation of $z$ and $d$ are two sequential linear transformations
  these can be combined into a single linear transformation.
  However we've implemented this separately so that it matches with the description
  in the paper.
docstr_tok: ["`","input_size","`","is","the","size","of","the","input","$","x_t","$",",","`","hidden_size","`","is","the","size","of","the","LSTM",",","and","`","hyper_size","`","is","the","size","of","the","smaller","LSTM","that","alters","the","weights","of","the","larger","outer","LSTM",".","`","n_z","`","is","the","size","of","the","feature","vectors","used","to","alter","the","LSTM","weights",".","We","use","the","output","of","the","smaller","LSTM","to","compute","$","z_h^","{","i",",","f",",","g",",","o","}","$",",","$","z_x^","{","i",",","f",",","g",",","o","}","$","and","$","z_b^","{","i",",","f",",","g",",","o","}","$","using","linear","transformations",".","We","calculate","$","d_h^","{","i",",","f",",","g",",","o","}","(","z_h^","{","i",",","f",",","g",",","o","}",")","$",",","$","d_x^","{","i",",","f",",","g",",","o","}","(","z_x^","{","i",",","f",",","g",",","o","}",")","$",",","and","$","d_b^","{","i",",","f",",","g",",","o","}","(","z_b^","{","i",",","f",",","g",",","o","}",")","$","from","these",",","using","linear","transformations","again",".","These","are","then","used","to","scale","the","rows","of","weight","and","bias","tensors","of","the","main","LSTM",".","\ud83d\udcdd","Since","the","computation","of","$","z","$","and","$","d","$","are","two","sequential","linear","transformations","these","can","be","combined","into","a","single","linear","transformation",".","However","we","'ve","implemented","this","separately","so","that","it","matches","with","the","description","in","the","paper","."]
code: def __init__(
      self, input_size: int, hidden_size: int, hyper_size: int, n_z: int
  ):
      """
      `input_size` is the size of the input $x_t$,
      `hidden_size` is the size of the LSTM, and
      `hyper_size` is the size of the smaller LSTM that alters the weights of the larger outer LSTM.
      `n_z` is the size of the feature vectors used to alter the LSTM weights.
  
      We use the output of the smaller LSTM to compute $z_h^{i,f,g,o}$, $z_x^{i,f,g,o}$ and
      $z_b^{i,f,g,o}$ using linear transformations.
      We calculate $d_h^{i,f,g,o}(z_h^{i,f,g,o})$, $d_x^{i,f,g,o}(z_x^{i,f,g,o})$, and
      $d_b^{i,f,g,o}(z_b^{i,f,g,o})$ from these, using linear transformations again.
      These are then used to scale the rows of weight and bias tensors of the main LSTM.
  
      📝 Since the computation of $z$ and $d$ are two sequential linear transformations
      these can be combined into a single linear transformation.
      However we've implemented this separately so that it matches with the description
      in the paper.
      """
      super().__init__()
  
      # The input to the hyperLSTM is
      # $$
      # \hat{x}_t = \begin{pmatrix}
      # h_{t-1} \\
      # x_t
      # \end{pmatrix}
      # $$
      # where $x_t$ is the input and $h_{t-1}$ is the output of the outer LSTM at previous step.
      # So the input size is `hidden_size + input_size`.
      #
      # The output of hyperLSTM is $\hat{h}_t$ and $\hat{c}_t$.
      self.hyper = LSTMCell(
          hidden_size + input_size, hyper_size, layer_norm=True
      )
  
      # $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_t)$$
      # 🤔 In the paper it was specified as
      # $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_{\textcolor{red}{t-1}})$$
      # I feel that it's a typo.
      self.z_h = nn.Linear(hyper_size, 4 * n_z)
      # $$z_x^{i,f,g,o} = lin_x^{i,f,g,o}(\hat{h}_t)$$
      self.z_x = nn.Linear(hyper_size, 4 * n_z)
      # $$z_b^{i,f,g,o} = lin_b^{i,f,g,o}(\hat{h}_t)$$
      self.z_b = nn.Linear(hyper_size, 4 * n_z, bias=False)
  
      # $$d_h^{i,f,g,o}(z_h^{i,f,g,o}) = lin_{dh}^{i,f,g,o}(z_h^{i,f,g,o})$$
      d_h = [nn.Linear(n_z, hidden_size, bias=False) for _ in range(4)]
      self.d_h = nn.ModuleList(d_h)
      # $$d_x^{i,f,g,o}(z_x^{i,f,g,o}) = lin_{dx}^{i,f,g,o}(z_x^{i,f,g,o})$$
      d_x = [nn.Linear(n_z, hidden_size, bias=False) for _ in range(4)]
      self.d_x = nn.ModuleList(d_x)
      # $$d_b^{i,f,g,o}(z_b^{i,f,g,o}) = lin_{db}^{i,f,g,o}(z_b^{i,f,g,o})$$
      d_b = [nn.Linear(n_z, hidden_size) for _ in range(4)]
      self.d_b = nn.ModuleList(d_b)
  
      # The weight matrices $W_h^{i,f,g,o}$
      self.w_h = nn.ParameterList(
          [
              nn.Parameter(torch.zeros(hidden_size, hidden_size))
              for _ in range(4)
          ]
      )
      # The weight matrices $W_x^{i,f,g,o}$
      self.w_x = nn.ParameterList(
          [
              nn.Parameter(torch.zeros(hidden_size, input_size))
              for _ in range(4)
          ]
      )
  
      # Layer normalization
      self.layer_norm = nn.ModuleList(
          [nn.LayerNorm(hidden_size) for _ in range(4)]
      )
      self.layer_norm_c = nn.LayerNorm(hidden_size)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/hypernetworks/hyper_lstm.py
imports: from typing import Optional, Tuple
  import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.lstm import LSTMCell
label_desc: Use image node_id 1 to create a new HyperLSTMCell object from inherited base classes: Module with example: obj = HyperLSTMCell(input_size, hidden_size, hyper_size, n_z)