class_docstr: ## Regularization loss
  
  $$L_{Reg} = \mathop{KL} \Big(p_n \Vert p_G(\lambda_p) \Big)$$
  
  $\mathop{KL}$ is the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
  
  $p_G$ is the [Geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) parameterized by
  $\lambda_p$. *$\lambda_p$ has nothing to do with $\lambda_n$; we are just sticking to same notation as the paper*.
  $$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$.
  
  The regularization loss biases the network towards taking $rac{1}{\lambda_p}$ steps and incentivies non-zero probabilities
  for all steps; i.e. promotes exploration.
class_docstr_tok: ["#","#","Regularization","loss","$","$","L_","{","Reg","}","=","\\mathop","{","KL","}","\\Big","(","p_n","\\Vert","p_G","(","\\lambda_p",")","\\Big",")","$","$","$","\\mathop","{","KL","}","$","is","the","[","Kullback\u2013Leibler","divergence","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Kullback","%","E2","%","80","%","93Leibler_divergence",")",".","$","p_G","$","is","the","[","Geometric","distribution","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Geometric_distribution",")","parameterized","by","$","\\lambda_p","$",".","*","$","\\lambda_p","$","has","nothing","to","do","with","$","\\lambda_n","$",";","we","are","just","sticking","to","same","notation","as","the","paper","*",".","$","$","Pr_","{","p_G","(","\\lambda_p",")","}","(","X","=","k",")","=","(","1","-","\\lambda_p",")","^k","\\lambda_p","$","$",".","The","regularization","loss","biases","the","network","towards","taking","$","rac","{","1","}","{","\\lambda_p","}","$","steps","and","incentivies","non-zero","probabilities","for","all","steps",";","i.e",".","promotes","exploration","."]
docstr: * `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
docstr_tok: ["*","`","p","`","is","$","p_1","\\dots","p_N","$","in","a","tensor","of","shape","`","[","N",",","batch_size","]","`"]
code: def forward(self, p: torch.Tensor):
      """
      * `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
      """
      # Transpose `p` to `[batch_size, N]`
      p = p.transpose(0, 1)
      # Get $Pr_{p_G(\lambda_p)}$ upto $N$ and expand it across the batch dimension
      p_g = self.p_g[None, : p.shape[1]].expand_as(p)
  
      # Calculate the KL-divergence.
      # *The [PyTorch KL-divergence](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html)
      # implementation accepts log probabilities.*
      return self.kl_div(p.log(), p_g)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/adaptive_computation/ponder_net/__init__.py
imports: from typing import Tuple
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the RegularizationLoss obj's underlying member method code with example usage: obj.forward(p) and returns: self