class_docstr: ## Hourglass model
  
  This model recursively adds layers to the middle while shortening the sequence by down-sampling.
  The shortened sequence processed by another hourglass model is sandwiched between two normal transformer
  layers. (A transformer layer has a [self-attention layer](../mha.html)
   and a [position-wise feed-forward layer](../feed_forward.html)).
docstr: * `n_heads` is the number of heads in [multi-head attention layers](../mha.html)
  * `d_model` is the size of the token embeddings
  * `dropout` is the dropout probability
  * `d_ff` is the dimensionality of the hidden layer in [position-wise feed-forward layers](../feed_forward.html)
  * `shortening_factors` is the list of shortening factors
code: def __init__(
      self,
      n_heads: int,
      d_model: int,
      dropout: float,
      d_ff: int,
      shortening_factors: List[int],
  ):
      """
      * `n_heads` is the number of heads in [multi-head attention layers](../mha.html)
      * `d_model` is the size of the token embeddings
      * `dropout` is the dropout probability
      * `d_ff` is the dimensionality of the hidden layer in [position-wise feed-forward layers](../feed_forward.html)
      * `shortening_factors` is the list of shortening factors
      """
      super().__init__()
  
      # The transformer layer before down-sampling
      self.pre = TransformerLayer(
          d_model=d_model,
          # [Multi-head attention layer](../mha.html)
          self_attn=MultiHeadAttention(n_heads, d_model, dropout),
          # [Position wise feed-forward layers](.. / feed_forward.html)
          feed_forward=FeedForward(d_model, d_ff, dropout),
          #
          dropout_prob=dropout,
      )
      # Auto-regressive mask
      self.mask = AutoregressiveMask()
  
      # The shortening factor $k$ (or the down-sampling rate)
      k = shortening_factors[0]
  
      # We shift the tokens to the right by $k - 1$ steps to make sure
      # information doesn't leak from the future tokens to past tokens
      # as a result of down-sampling and up-sampling
      self.shift_right = ShiftRight(k - 1)
      # Shortening or the down-sampling layer. We use the simplest form - average pooling.
      # The paper shows that attention based down sampling works best, which we haven't implemented yet.
      self.shortening = AvgPoolShortening(k)
  
      # If there are no more shortening (middle of the hourglass)
      if len(shortening_factors) == 1:
          # The center layer is another transformer layer
          self.shortened = TransformerLayer(
              d_model=d_model,
              self_attn=MultiHeadAttention(n_heads, d_model, dropout),
              feed_forward=FeedForward(d_model, d_ff, dropout),
              dropout_prob=dropout,
          )
          # Autoregressive mask
          self.mask_short = AutoregressiveMask()
          self.hour_glass = None
      else:
          # Insert another hourglass model recursively
          self.hour_glass = HourGlass(
              n_heads, d_model, dropout, d_ff, shortening_factors[1:]
          )
  
      # Up-sampling layer. We use naive up-sampling for simplicity and the paper shows attention based up sampling
      # works better.
      self.up_sampling = NaiveUpSampling(k)
  
      # The final transformer layer after up-sampling
      self.post = TransformerLayer(
          d_model=d_model,
          self_attn=MultiHeadAttention(n_heads, d_model, dropout),
          feed_forward=FeedForward(d_model, d_ff, dropout),
          dropout_prob=dropout,
      )
  
imports: from typing import List
  import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers import MultiHeadAttention, TransformerLayer
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.utils import subsequent_mask
label_desc: Use image node_id 1 to create a new HourGlass object from inherited base classes: Module with example: obj = HourGlass(n_heads, d_model, dropout, d_ff, shortening_factors)