class_docstr: <a id="KLDivergenceLoss"></a>
  
  ## KL Divergence Regularization Loss
  
  This tries to shrink the total evidence to zero if the sample cannot be correctly classified.
  
  First we calculate $        ilde{lpha}_k = y_k + (1 - y_k)         extcolor{orange}{lpha_k}$ the
  Dirichlet parameters after remove the correct evidence.
  
  egin{align}
  &KL \Big[ D(\mathbf{p} ert \mathbf{        ilde{lpha}}) \Big \Vert
  D(\mathbf{p} ert <1, \dots, 1>\Big] \
  &= \log \Bigg( rac{\Gamma \Big( \sum_{k=1}^K       ilde{lpha}_k \Big)}
  {\Gamma(K) \prod_{k=1}^K \Gamma(    ilde{lpha}_k)} \Bigg)
  + \sum_{k=1}^K (    ilde{lpha}_k - 1)
  \Big[ \psi( ilde{lpha}_k) - \psi(  ilde{S}) \Big]
  \end{align}
  
  where $\Gamma(\cdot)$ is the gamma function,
  $\psi(\cdot)$ is the $digamma$ function and
  $   ilde{S} = \sum_{k=1}^K  ilde{lpha}_k$
class_docstr_tok: ["<","a","id=","''","KLDivergenceLoss","''",">","<","\/a",">","#","#","KL","Divergence","Regularization","Loss","This","tries","to","shrink","the","total","evidence","to","zero","if","the","sample","can","not","be","correctly","classified",".","First","we","calculate","$","ilde","{","\u0007lpha","}","_k","=","y_k","+","(","1","-","y_k",")","extcolor","{","orange","}","{","\u0007lpha_k","}","$","the","Dirichlet","parameters","after","remove","the","correct","evidence",".","\begin","{","align","}","&","KL","\\Big","[","D","(","\\mathbf","{","p","}","ert","\\mathbf","{","ilde","{","\u0007lpha","}","}",")","\\Big","\\Vert","D","(","\\mathbf","{","p","}","ert","<","1",",","\\dots",",","1",">","\\Big","]","\\","&","=","\\log","\\Bigg","(","rac","{","\\Gamma","\\Big","(","\\sum_","{","k=1","}","^K","ilde","{","\u0007lpha","}","_k","\\Big",")","}","{","\\Gamma","(","K",")","\\prod_","{","k=1","}","^K","\\Gamma","(","ilde","{","\u0007lpha","}","_k",")","}","\\Bigg",")","+","\\sum_","{","k=1","}","^K","(","ilde","{","\u0007lpha","}","_k","-","1",")","\\Big","[","\\psi","(","ilde","{","\u0007lpha","}","_k",")","-","\\psi","(","ilde","{","S","}",")","\\Big","]","\\end","{","align","}","where","$","\\Gamma","(","\\cdot",")","$","is","the","gamma","function",",","$","\\psi","(","\\cdot",")","$","is","the","$","digamma","$","function","and","$","ilde","{","S","}","=","\\sum_","{","k=1","}","^K","ilde","{","\u0007lpha","}","_k","$"]
docstr: * `evidence` is $\mathbf{e} \ge 0$ with shape `[batch_size, n_classes]`
  * `target` is $\mathbf{y}$ with shape `[batch_size, n_classes]`
docstr_tok: ["*","`","evidence","`","is","$","\\mathbf","{","e","}","\\ge","0","$","with","shape","`","[","batch_size",",","n_classes","]","`","*","`","target","`","is","$","\\mathbf","{","y","}","$","with","shape","`","[","batch_size",",","n_classes","]","`"]
code: def forward(self, evidence: torch.Tensor, target: torch.Tensor):
      """
      * `evidence` is $\mathbf{e} \ge 0$ with shape `[batch_size, n_classes]`
      * `target` is $\mathbf{y}$ with shape `[batch_size, n_classes]`
      """
      # $\textcolor{orange}{\alpha_k} = e_k + 1$
      alpha = evidence + 1.0
      # Number of classes
      n_classes = evidence.shape[-1]
      # Remove non-misleading evidence
      # $$\tilde{\alpha}_k = y_k + (1 - y_k) \textcolor{orange}{\alpha_k}$$
      alpha_tilde = target + (1 - target) * alpha
      # $\tilde{S} = \sum_{k=1}^K \tilde{\alpha}_k$
      strength_tilde = alpha_tilde.sum(dim=-1)
  
      # The first term
      #
      # \begin{align}
      # &\log \Bigg( \frac{\Gamma \Big( \sum_{k=1}^K \tilde{\alpha}_k \Big)}
      #     {\Gamma(K) \prod_{k=1}^K \Gamma(\tilde{\alpha}_k)} \Bigg) \\
      # &= \log \Gamma \Big( \sum_{k=1}^K \tilde{\alpha}_k \Big)
      #   - \log \Gamma(K)
      #   - \sum_{k=1}^K \log \Gamma(\tilde{\alpha}_k)
      # \end{align}
      first = (
          torch.lgamma(alpha_tilde.sum(dim=-1))
          - torch.lgamma(alpha_tilde.new_tensor(float(n_classes)))
          - (torch.lgamma(alpha_tilde)).sum(dim=-1)
      )
  
      # The second term
      # $$\sum_{k=1}^K (\tilde{\alpha}_k - 1)
      #     \Big[ \psi(\tilde{\alpha}_k) - \psi(\tilde{S}) \Big]$$
      second = (
          (alpha_tilde - 1)
          * (
              torch.digamma(alpha_tilde)
              - torch.digamma(strength_tilde)[:, None]
          )
      ).sum(dim=-1)
  
      # Sum of the terms
      loss = first + second
  
      # Mean loss over the batch
      return loss.mean()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/uncertainty/evidence/__init__.py
imports: import torch
  from labml import tracker
  from labml_helpers.module import Module
label_desc: Use image node_id 1 for calling the KLDivergenceLoss obj's underlying member method code with example usage: obj.forward(evidence, target) and returns: loss