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","."]
code: def forward(
      self,
      x: torch.Tensor,
      h: torch.Tensor,
      c: torch.Tensor,
      h_hat: torch.Tensor,
      c_hat: torch.Tensor,
  ):
      # $$
      # \hat{x}_t = \begin{pmatrix}
      # h_{t-1} \\
      # x_t
      # \end{pmatrix}
      # $$
      x_hat = torch.cat((h, x), dim=-1)
      # $$\hat{h}_t, \hat{c}_t = lstm(\hat{x}_t, \hat{h}_{t-1}, \hat{c}_{t-1})$$
      h_hat, c_hat = self.hyper(x_hat, h_hat, c_hat)
  
      # $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_t)$$
      z_h = self.z_h(h_hat).chunk(4, dim=-1)
      # $$z_x^{i,f,g,o} = lin_x^{i,f,g,o}(\hat{h}_t)$$
      z_x = self.z_x(h_hat).chunk(4, dim=-1)
      # $$z_b^{i,f,g,o} = lin_b^{i,f,g,o}(\hat{h}_t)$$
      z_b = self.z_b(h_hat).chunk(4, dim=-1)
  
      # We calculate $i$, $f$, $g$ and $o$ in a loop
      ifgo = []
      for i in range(4):
          # $$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 = self.d_h[i](z_h[i])
          # $$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 = self.d_x[i](z_x[i])
  
          # \begin{align}
          # {i,f,g,o} = LN(&\textcolor{lightgreen}{d_h^{i,f,g,o}(z_h) \odot (W_h^{i,f,g,o} h_{t-1})} \\
          #              + &\textcolor{lightgreen}{d_x^{i,f,g,o}(z_x) \odot (W_h^{i,f,g,o} x_t)} \\
          #              + &d_b^{i,f,g,o}(z_b))
          # \end{align}
          y = (
              d_h * torch.einsum("ij,bj->bi", self.w_h[i], h)
              + d_x * torch.einsum("ij,bj->bi", self.w_x[i], x)
              + self.d_b[i](z_b[i])
          )
  
          ifgo.append(self.layer_norm[i](y))
  
      # $$i_t, f_t, g_t, o_t$$
      i, f, g, o = ifgo
  
      # $$c_t = \sigma(f_t) \odot c_{t-1} + \sigma(i_t) \odot \tanh(g_t) $$
      c_next = torch.sigmoid(f) * c + torch.sigmoid(i) * torch.tanh(g)
  
      # $$h_t = \sigma(o_t) \odot \tanh(LN(c_t))$$
      h_next = torch.sigmoid(o) * torch.tanh(self.layer_norm_c(c_next))
  
      return h_next, c_next, h_hat, c_hat
  
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 2 for calling the HyperLSTMCell obj's underlying member method code with example usage: obj.forward(x, h, c, h_hat, c_hat) and returns: h_next,  c_next,  h_hat,  c_hat