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: * `n_elems` is the number of elements in the input vector
  * `n_hidden` is the state vector size of the GRU
  * `max_steps` is the maximum number of steps $N$
docstr_tok: ["*","`","n_elems","`","is","the","number","of","elements","in","the","input","vector","*","`","n_hidden","`","is","the","state","vector","size","of","the","GRU","*","`","max_steps","`","is","the","maximum","number","of","steps","$","N","$"]
code: def __init__(self, n_elems: int, n_hidden: int, max_steps: int):
      """
      * `n_elems` is the number of elements in the input vector
      * `n_hidden` is the state vector size of the GRU
      * `max_steps` is the maximum number of steps $N$
      """
      super().__init__()
  
      self.max_steps = max_steps
      self.n_hidden = n_hidden
  
      # GRU
      # $$h_{n+1} = s_h(x, h_n)$$
      self.gru = nn.GRUCell(n_elems, n_hidden)
      # $$\hat{y}_n = s_y(h_n)$$
      # We could use a layer that takes the concatenation of $h$ and $x$ as input
      # but we went with this for simplicity.
      self.output_layer = nn.Linear(n_hidden, 1)
      # $$\lambda_n = s_\lambda(h_n)$$
      self.lambda_layer = nn.Linear(n_hidden, 1)
      self.lambda_prob = nn.Sigmoid()
      # An option to set during inference so that computation is actually halted at inference time
      self.is_halt = False
  
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 ParityPonderGRU object from inherited base classes: Module with example: obj = ParityPonderGRU(n_elems, n_hidden, max_steps)