class_docstr: ## Reconstruction Loss
class_docstr_tok: ["#","#","Reconstruction","Loss"]
code: def forward(
      self,
      mask: torch.Tensor,
      target: torch.Tensor,
      dist: "BivariateGaussianMixture",
      q_logits: torch.Tensor,
  ):
      # Get $\Pi$ and $\mathcal{N}(\mu_{x}, \mu_{y}, \sigma_{x}, \sigma_{y}, \rho_{xy})$
      pi, mix = dist.get_distribution()
      # `target` has shape `[seq_len, batch_size, 5]` where the last dimension is the features
      # $(\Delta x, \Delta y, p_1, p_2, p_3)$.
      # We want to get $\Delta x, \Delta$ y and get the probabilities from each of the distributions
      # in the mixture $\mathcal{N}(\mu_{x}, \mu_{y}, \sigma_{x}, \sigma_{y}, \rho_{xy})$.
      #
      # `xy` will have shape `[seq_len, batch_size, n_distributions, 2]`
      xy = (
          target[:, :, 0:2]
          .unsqueeze(-2)
          .expand(-1, -1, dist.n_distributions, -1)
      )
      # Calculate the probabilities
      # $$p(\Delta x, \Delta y) =
      # \sum_{j=1}^M \Pi_j \mathcal{N} \big( \Delta x, \Delta y \vert
      # \mu_{x,j}, \mu_{y,j}, \sigma_{x,j}, \sigma_{y,j}, \rho_{xy,j}
      # \big)$$
      probs = torch.sum(pi.probs * torch.exp(mix.log_prob(xy)), 2)
  
      # $$L_s = - \frac{1}{N_{max}} \sum_{i=1}^{N_s} \log \big (p(\Delta x, \Delta y) \big)$$
      # Although `probs` has $N_{max}$ (`longest_seq_len`) elements, the sum is only taken
      # upto $N_s$ because the rest is masked out.
      #
      # It might feel like we should be taking the sum and dividing by $N_s$ and not $N_{max}$,
      # but this will give higher weight for individual predictions in shorter sequences.
      # We give equal weight to each prediction $p(\Delta x, \Delta y)$ when we divide by $N_{max}$
      loss_stroke = -torch.mean(mask * torch.log(1e-5 + probs))
  
      # $$L_p = - \frac{1}{N_{max}} \sum_{i=1}^{N_{max}} \sum_{k=1}^{3} p_{k,i} \log(q_{k,i})$$
      loss_pen = -torch.mean(target[:, :, 2:] * q_logits)
  
      # $$L_R = L_s + L_p$$
      return loss_stroke + loss_pen
  
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 for calling the ReconstructionLoss obj's underlying member method code with example usage: obj.forward(mask, target, dist, q_logits) and returns: unknown