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"]
code: def initial_state(self, batch_size: Optional[int]):
      state = self._initial_state
      if batch_size is not None:
          state = add_batch(state, batch_size)
      return state
  
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 3 for calling the IdentityRNN obj's underlying member method code with example usage: obj.initial_state(batch_size) and returns: state