class_docstr: Basic fully-connected RNN core with identity initialization.
  
  Given :math:`x_t` and the previous hidden state :math:`h_{t-1}` the
  core computes
  .. math::
     h_t = \operatorname{ReLU}(w_i x_t + b_i + w_h h_{t-1} + b_h)
  The output is equal to the new state, :math:`h_t`.
  
  Initialized using the strategy described in:
    https://arxiv.org/pdf/1504.00941.pdf
class_docstr_tok: ["Basic","fully-connected","RNN","core","with","identity","initialization",".","Given",":","math",":","`","x_t","`","and","the","previous","hidden","state",":","math",":","`","h_","{","t-1","}","`","the","core","computes","..","math",":",":","h_t","=","\\operatorname","{","ReLU","}","(","w_i","x_t","+","b_i","+","w_h","h_","{","t-1","}","+","b_h",")","The","output","is","equal","to","the","new","state",",",":","math",":","`","h_t","`",".","Initialized","using","the","strategy","described","in",":","https",":","\/\/arxiv.org\/pdf\/1504.00941.pdf"]
docstr: Constructs a vanilla RNN core.
  
  Args:
    hidden_size: Hidden layer size.
    hidden_scale: Scalar multiplying the hidden-to-hidden matmul.
    name: Name of the module.
docstr_tok: ["Constructs","a","vanilla","RNN","core",".","Args",":","hidden_size",":","Hidden","layer","size",".","hidden_scale",":","Scalar","multiplying","the","hidden-to-hidden","matmul",".","name",":","Name","of","the","module","."]
code: def __init__(
      self,
      hidden_size: int,
      hidden_scale: float = 1e-2,
      name: Optional[str] = None,
  ):
      """Constructs a vanilla RNN core.
  
      Args:
        hidden_size: Hidden layer size.
        hidden_scale: Scalar multiplying the hidden-to-hidden matmul.
        name: Name of the module.
      """
      super().__init__(name=name)
      self._initial_state = jnp.zeros([hidden_size])
      self._hidden_scale = hidden_scale
      self._input_to_hidden = hk.Linear(hidden_size)
      self._hidden_to_hidden = hk.Linear(
          hidden_size, with_bias=True, w_init=w_init_identity
      )
  
imports: import dataclasses
  from typing import Callable, NamedTuple, Optional, Sequence, Tuple, Union
  from acme import specs
  from acme.agents.jax.mpo import types
  from acme.jax import networks
  from acme.jax import utils
  import chex
  import haiku
  import haiku.initializers
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 1 to create a new IdentityRNN object from inherited base classes: hk with example: obj = IdentityRNN(hidden_size, hidden_scale, name)