class_docstr: ## Multilayer LSTM
class_docstr_tok: ["#","#","Multilayer","LSTM"]
docstr: `x` has shape `[n_steps, batch_size, input_size]` and
  `state` is a tuple of $h$ and $c$, each with a shape of `[batch_size, hidden_size]`.
docstr_tok: ["`","x","`","has","shape","`","[","n_steps",",","batch_size",",","input_size","]","`","and","`","state","`","is","a","tuple","of","$","h","$","and","$","c","$",",","each","with","a","shape","of","`","[","batch_size",",","hidden_size","]","`","."]
code: def forward(
      self,
      x: torch.Tensor,
      state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  ):
      """
      `x` has shape `[n_steps, batch_size, input_size]` and
      `state` is a tuple of $h$ and $c$, each with a shape of `[batch_size, hidden_size]`.
      """
      n_steps, batch_size = x.shape[:2]
  
      # Initialize the state if `None`
      if state is None:
          h = [
              x.new_zeros(batch_size, self.hidden_size)
              for _ in range(self.n_layers)
          ]
          c = [
              x.new_zeros(batch_size, self.hidden_size)
              for _ in range(self.n_layers)
          ]
      else:
          (h, c) = state
          # Reverse stack the tensors to get the states of each layer
          #
          # 📝 You can just work with the tensor itself but this is easier to debug
          h, c = list(torch.unbind(h)), list(torch.unbind(c))
  
      # Array to collect the outputs of the final layer at each time step.
      out = []
      for t in range(n_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
              h[layer], c[layer] = self.cells[layer](
                  inp, h[layer], c[layer]
              )
              # Input to the next layer is the state of this layer
              inp = h[layer]
          # Collect the output $h$ of the final layer
          out.append(h[-1])
  
      # Stack the outputs and states
      out = torch.stack(out)
      h = torch.stack(h)
      c = torch.stack(c)
  
      return out, (h, c)
  
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 2 for calling the LSTM obj's underlying member method code with example usage: obj.forward(x, state) and returns: out,  h,  c