class_docstr: ## PonderNet with GRU for Parity Task
  
  This is a simple model that uses a [GRU Cell](https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html)
  as the step function.
  
  This model is for the [Parity Task](../parity.html) where the input is a vector of `n_elems`.
  Each element of the vector is either `0`, `1` or `-1` and the output is the parity
  - a binary value that is true if the number of `1`s is odd and false otherwise.
  
  The prediction of the model is the log probability of the parity being $1$.
class_docstr_tok: ["#","#","PonderNet","with","GRU","for","Parity","Task","This","is","a","simple","model","that","uses","a","[","GRU","Cell","]","(","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.nn.GRUCell.html",")","as","the","step","function",".","This","model","is","for","the","[","Parity","Task","]","(","..","\/parity.html",")","where","the","input","is","a","vector","of","`","n_elems","`",".","Each","element","of","the","vector","is","either","`","0","`",",","`","1","`","or","`","-1","`","and","the","output","is","the","parity","-","a","binary","value","that","is","true","if","the","number","of","`","1","`","s","is","odd","and","false","otherwise",".","The","prediction","of","the","model","is","the","log","probability","of","the","parity","being","$","1","$","."]
docstr: * `x` is the input of shape `[batch_size, n_elems]`
  
  This outputs a tuple of four tensors:
  
  1. $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
  2. $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size]` - the log probabilities of the parity being $1$
  3. $p_m$ of shape `[batch_size]`
  4. $\hat{y}_m$ of shape `[batch_size]` where the computation was halted at step $m$
docstr_tok: ["*","`","x","`","is","the","input","of","shape","`","[","batch_size",",","n_elems","]","`","This","outputs","a","tuple","of","four","tensors",":","1",".","$","p_1","\\dots","p_N","$","in","a","tensor","of","shape","`","[","N",",","batch_size","]","`","2",".","$","\\hat","{","y","}","_1","\\dots","\\hat","{","y","}","_N","$","in","a","tensor","of","shape","`","[","N",",","batch_size","]","`","-","the","log","probabilities","of","the","parity","being","$","1","$","3",".","$","p_m","$","of","shape","`","[","batch_size","]","`","4",".","$","\\hat","{","y","}","_m","$","of","shape","`","[","batch_size","]","`","where","the","computation","was","halted","at","step","$","m","$"]
code: def forward(
      self, x: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
      """
      * `x` is the input of shape `[batch_size, n_elems]`
  
      This outputs a tuple of four tensors:
  
      1. $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
      2. $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size]` - the log probabilities of the parity being $1$
      3. $p_m$ of shape `[batch_size]`
      4. $\hat{y}_m$ of shape `[batch_size]` where the computation was halted at step $m$
      """
  
      #
      batch_size = x.shape[0]
  
      # We get initial state $h_1 = s_h(x)$
      h = x.new_zeros((x.shape[0], self.n_hidden))
      h = self.gru(x, h)
  
      # Lists to store $p_1 \dots p_N$ and $\hat{y}_1 \dots \hat{y}_N$
      p = []
      y = []
      # $\prod_{j=1}^{n-1} (1 - \lambda_j)$
      un_halted_prob = h.new_ones((batch_size,))
  
      # A vector to maintain which samples has halted computation
      halted = h.new_zeros((batch_size,))
      # $p_m$ and $\hat{y}_m$ where the computation was halted at step $m$
      p_m = h.new_zeros((batch_size,))
      y_m = h.new_zeros((batch_size,))
  
      # Iterate for $N$ steps
      for n in range(1, self.max_steps + 1):
          # The halting probability $\lambda_N = 1$ for the last step
          if n == self.max_steps:
              lambda_n = h.new_ones(h.shape[0])
          # $\lambda_n = s_\lambda(h_n)$
          else:
              lambda_n = self.lambda_prob(self.lambda_layer(h))[:, 0]
          # $\hat{y}_n = s_y(h_n)$
          y_n = self.output_layer(h)[:, 0]
  
          # $$p_n = \lambda_n \prod_{j=1}^{n-1} (1 - \lambda_j)$$
          p_n = un_halted_prob * lambda_n
          # Update $\prod_{j=1}^{n-1} (1 - \lambda_j)$
          un_halted_prob = un_halted_prob * (1 - lambda_n)
  
          # Halt based on halting probability $\lambda_n$
          halt = torch.bernoulli(lambda_n) * (1 - halted)
  
          # Collect $p_n$ and $\hat{y}_n$
          p.append(p_n)
          y.append(y_n)
  
          # Update $p_m$ and $\hat{y}_m$ based on what was halted at current step $n$
          p_m = p_m * (1 - halt) + p_n * halt
          y_m = y_m * (1 - halt) + y_n * halt
  
          # Update halted samples
          halted = halted + halt
          # Get next state $h_{n+1} = s_h(x, h_n)$
          h = self.gru(x, h)
  
          # Stop the computation if all samples have halted
          if self.is_halt and halted.sum() == batch_size:
              break
  
      #
      return torch.stack(p), torch.stack(y), p_m, y_m
  
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 ParityPonderGRU obj's underlying member method code with example usage: obj.forward(x) and returns: torch,  torch,  p_m,  y_m