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: * `lambda_p` is $\lambda_p$ - the success probability of geometric distribution
  * `max_steps` is the highest $N$; we use this to pre-compute $p_G(\lambda_p)$
docstr_tok: ["*","`","lambda_p","`","is","$","\\lambda_p","$","-","the","success","probability","of","geometric","distribution","*","`","max_steps","`","is","the","highest","$","N","$",";","we","use","this","to","pre-compute","$","p_G","(","\\lambda_p",")","$"]
code: def __init__(self, lambda_p: float, max_steps: int = 1_000):
      """
      * `lambda_p` is $\lambda_p$ - the success probability of geometric distribution
      * `max_steps` is the highest $N$; we use this to pre-compute $p_G(\lambda_p)$
      """
      super().__init__()
  
      # Empty vector to calculate $p_G(\lambda_p)$
      p_g = torch.zeros((max_steps,))
      # $(1 - \lambda_p)^k$
      not_halted = 1.0
      # Iterate upto `max_steps`
      for k in range(max_steps):
          # $$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$
          p_g[k] = not_halted * lambda_p
          # Update $(1 - \lambda_p)^k$
          not_halted = not_halted * (1 - lambda_p)
  
      # Save $Pr_{p_G(\lambda_p)}$
      self.p_g = nn.Parameter(p_g, requires_grad=False)
  
      # KL-divergence loss
      self.kl_div = nn.KLDivLoss(reduction="batchmean")
  
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 1 to create a new RegularizationLoss object from inherited base classes: Module with example: obj = RegularizationLoss(lambda_p, max_steps)