class_docstr: ## Multilayer Recurrent Highway Network
class_docstr_tok: ["#","#","Multilayer","Recurrent","Highway","Network"]
docstr: `x` has shape `[seq_len, batch_size, input_size]` and
  `state` has shape `[batch_size, hidden_size]`.
docstr_tok: ["`","x","`","has","shape","`","[","seq_len",",","batch_size",",","input_size","]","`","and","`","state","`","has","shape","`","[","batch_size",",","hidden_size","]","`","."]
code: def forward(
      self, x: torch.Tensor, state: Optional[torch.Tensor] = None
  ):
      """
      `x` has shape `[seq_len, batch_size, input_size]` and
      `state` has shape `[batch_size, hidden_size]`.
      """
      time_steps, batch_size = x.shape[:2]
  
      # Initialize the state if `None`
      if state is None:
          s = [
              x.new_zeros(batch_size, self.hidden_size)
              for _ in range(self.n_layers)
          ]
      else:
          # Reverse stack the state to get the state of each layer
          #
          # 📝 You can just work with the tensor itself but this is easier to debug
          s = torch.unbind(state)
  
      # Array to collect the outputs of the final layer at each time step.
      out = []
  
      # Run through the network for each time step
      for t in range(time_steps):
          # Input to the first layer is the input itself
          inp = x[t]
          # Loop through the layers
          for layer in range(self.n_layers):
              # Get the state of the layer
              s[layer] = self.cells[layer](inp, s[layer])
              # Input to the next layer is the state of this layer
              inp = s[layer]
          # Collect the output of the final layer
          out.append(s[-1])
  
      # Stack the outputs and states
      out = torch.stack(out)
      s = torch.stack(s)
  
      return out, 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 RHN obj's underlying member method code with example usage: obj.forward(x, state) and returns: out,  s