class_docstr: ## Decoder module
  
  This consists of a LSTM
class_docstr_tok: ["#","#","Decoder","module","This","consists","of","a","LSTM"]
code: def __init__(
      self, d_z: int, dec_hidden_size: int, n_distributions: int
  ):
      super().__init__()
      # LSTM takes $[(\Delta x, \Delta y, p_1, p_2, p_3); z]$ as input
      self.lstm = nn.LSTM(d_z + 5, dec_hidden_size)
  
      # Initial state of the LSTM is $[h_0; c_0] = \tanh(W_{z}z + b_z)$.
      # `init_state` is the linear transformation for this
      self.init_state = nn.Linear(d_z, 2 * dec_hidden_size)
  
      # This layer produces outputs for each of the `n_distributions`.
      # Each distribution needs six parameters
      # $(\hat{\Pi_i}, \mu_{x_i}, \mu_{y_i}, \hat{\sigma_{x_i}}, \hat{\sigma_{y_i}} \hat{\rho_{xy_i}})$
      self.mixtures = nn.Linear(dec_hidden_size, 6 * n_distributions)
  
      # This head is for the logits $(\hat{q_1}, \hat{q_2}, \hat{q_3})$
      self.q_head = nn.Linear(dec_hidden_size, 3)
      # This is to calculate $\log(q_k)$ where
      # $$q_k = \operatorname{softmax}(\hat{q})_k = \frac{\exp(\hat{q_k})}{\sum_{j = 1}^3 \exp(\hat{q_j})}$$
      self.q_log_softmax = nn.LogSoftmax(-1)
  
      # These parameters are stored for future reference
      self.n_distributions = n_distributions
      self.dec_hidden_size = dec_hidden_size
  
imports: import math
  from typing import Optional, Tuple, Any
  import numpy
  import torch
  import torch.nn
  from matplotlib import pyplot
  from torch import optim
  from torch.utils.data import Dataset, DataLoader
  import einops
  from labml import lab, experiment, tracker, monit
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.module import Module
  from labml_helpers.optimizer import OptimizerConfigs
  from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
label_desc: Use image node_id 1 to create a new DecoderRNN object from inherited base classes: Module with example: obj = DecoderRNN(d_z, dec_hidden_size, n_distributions)