class_docstr: Recurrent Neural Network (RNN) block.
docstr: Forward pass through RNN.
  
  Args:
      inputs: ``JaxArray`` with dimensions ``num_steps, batch_size, vocabulary_size``.
      only_return_final: return only the last output if ``True``, or all output otherwise.`
  
  Returns:
      Output tensor with dimensions ``num_steps * batch_size, vocabulary_size``.
code: def __call__(
      self, inputs: JaxArray, only_return_final=False
  ) -> JaxArray:
      """Forward pass through RNN.
  
      Args:
          inputs: ``JaxArray`` with dimensions ``num_steps, batch_size, vocabulary_size``.
          only_return_final: return only the last output if ``True``, or all output otherwise.`
  
      Returns:
          Output tensor with dimensions ``num_steps * batch_size, vocabulary_size``.
      """
      # Dimensions: num_steps, batch_size, vocab_size
      outputs = []
      for x in inputs:
          self.state.value = self.activation(
              jn.dot(x, self.w_xh.value)
              + jn.dot(self.state.value, self.w_hh.value)
              + self.b_h.value
          )
          y = self.output_layer(self.state.value)
          outputs.append(y)
      if only_return_final:
          return outputs[-1]
      return jn.concatenate(outputs, axis=0)
  
file_path: objax/objax/zoo/rnn.py
imports: from typing import Callable
  import jax.numpy
  from objax import Module
  from objax.nn import Linear
  from objax.nn.init import kaiming_normal
  from objax.typing import JaxArray
  from objax.variable import TrainVar, StateVar
label_desc: Use image node_id 3 for calling the RNN obj's underlying member method code with example usage: obj.__call__(inputs, only_return_final) and returns: jn, outputs