class_docstr: ## Recurrent Highway Network Cell
  
  This implements equations $(6) - (9)$.
  
  $s_d^t = h_d^t \odot g_d^t + s_{d - 1}^t \odot c_d^t$
  
  where
  
  egin{align}
  h_0^t &=    anh(lin_{hx}(x) + lin_{hs}(s_D^{t-1})) \
  g_0^t &= \sigma(lin_{gx}(x) + lin_{gs}^1(s_D^{t-1})) \
  c_0^t &= \sigma(lin_{cx}(x) + lin_{cs}^1(s_D^{t-1}))
  \end{align}
  
  and for $0 < d < D$
  
  egin{align}
  h_d^t &=    anh(lin_{hs}^d(s_d^t)) \
  g_d^t &= \sigma(lin_{gs}^d(s_d^t)) \
  c_d^t &= \sigma(lin_{cs}^d(s_d^t))
  \end{align}
  
  $\odot$ stands for element-wise multiplication.
  
  Here we have made a couple of changes to notations from the paper.
  To avoid confusion with time, gate is represented with $g$,
  which was $t$ in the paper.
  To avoid confusion with multiple layers we use $d$ for depth and $D$ for
  total depth instead of $l$ and $L$ from the paper.
  
  We have also replaced the weight matrices and bias vectors from the equations with
  linear transforms, because that's how the implementation is going to look like.
  
  We implement weight tying, as described in paper, $c_d^t = 1 - g_d^t$.
class_docstr_tok: ["#","#","Recurrent","Highway","Network","Cell","This","implements","equations","$","(","6",")","-","(","9",")","$",".","$","s_d^t","=","h_d^t","\\odot","g_d^t","+","s_","{","d","-","1","}","^t","\\odot","c_d^t","$","where","\begin","{","align","}","h_0^t","&","=","anh","(","lin_","{","hx","}","(","x",")","+","lin_","{","hs","}","(","s_D^","{","t-1","}",")",")","\\","g_0^t","&","=","\\sigma","(","lin_","{","gx","}","(","x",")","+","lin_","{","gs","}","^1","(","s_D^","{","t-1","}",")",")","\\","c_0^t","&","=","\\sigma","(","lin_","{","cx","}","(","x",")","+","lin_","{","cs","}","^1","(","s_D^","{","t-1","}",")",")","\\end","{","align","}","and","for","$","0","<","d","<","D","$","\begin","{","align","}","h_d^t","&","=","anh","(","lin_","{","hs","}","^d","(","s_d^t",")",")","\\","g_d^t","&","=","\\sigma","(","lin_","{","gs","}","^d","(","s_d^t",")",")","\\","c_d^t","&","=","\\sigma","(","lin_","{","cs","}","^d","(","s_d^t",")",")","\\end","{","align","}","$","\\odot","$","stands","for","element-wise","multiplication",".","Here","we","have","made","a","couple","of","changes","to","notations","from","the","paper",".","To","avoid","confusion","with","time",",","gate","is","represented","with","$","g","$",",","which","was","$","t","$","in","the","paper",".","To","avoid","confusion","with","multiple","layers","we","use","$","d","$","for","depth","and","$","D","$","for","total","depth","instead","of","$","l","$","and","$","L","$","from","the","paper",".","We","have","also","replaced","the","weight","matrices","and","bias","vectors","from","the","equations","with","linear","transforms",",","because","that","'s","how","the","implementation","is","going","to","look","like",".","We","implement","weight","tying",",","as","described","in","paper",",","$","c_d^t","=","1","-","g_d^t","$","."]
docstr: `x` has shape `[batch_size, input_size]` and
  `s` has shape `[batch_size, hidden_size]`.
docstr_tok: ["`","x","`","has","shape","`","[","batch_size",",","input_size","]","`","and","`","s","`","has","shape","`","[","batch_size",",","hidden_size","]","`","."]
code: def forward(self, x: torch.Tensor, s: torch.Tensor):
      """
      `x` has shape `[batch_size, input_size]` and
      `s` has shape `[batch_size, hidden_size]`.
      """
  
      # Iterate $0 \leq d < D$
      for d in range(self.depth):
          # We calculate the concatenation of linear transforms for $h$ and $g$
          if d == 0:
              # The input is used only when $d$ is $0$.
              hg = self.input_lin(x) + self.hidden_lin[d](s)
          else:
              hg = self.hidden_lin[d](s)
  
          # Use the first half of `hg` to get $h_d^t$
          #
          # \begin{align}
          # h_0^t &= \tanh(lin_{hx}(x) + lin_{hs}(s_D^{t-1})) \\
          # h_d^t &= \tanh(lin_{hs}^d(s_d^t))
          # \end{align}
          h = torch.tanh(hg[:, : self.hidden_size])
          # Use the second half of `hg` to get $g_d^t$
          #
          # \begin{align}
          # g_0^t &= \sigma(lin_{gx}(x) + lin_{gs}^1(s_D^{t-1})) \\
          # g_d^t &= \sigma(lin_{gs}^d(s_d^t))
          # \end{align}
          g = torch.sigmoid(hg[:, self.hidden_size :])
  
          s = h * g + s * (1 - g)
  
      return s
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/recurrent_highway_networks/__init__.py
imports: from typing import Optional
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the RHNCell obj's underlying member method code with example usage: obj.forward(x, s) and returns: s