class_docstr: Builds a traditional GRU cell with dense internal transformations.
  
  Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555
class_docstr_tok: ["Builds","a","traditional","GRU","cell","with","dense","internal","transformations",".","Gated","Recurrent","Unit","paper",":","https",":","\/\/arxiv.org\/abs\/1412.3555"]
code: def forward(self, inputs):
      x, gru_state = inputs
  
      # Dense layer on the concatenation of x and h.
      w1, b1, w2, b2 = self.weights
      y = jnp.dot(jnp.concatenate([x, gru_state], axis=-1), w1) + b1
  
      # Update and reset gates.
      u, r = jnp.split(fastmath.sigmoid(y), 2, axis=-1)
  
      # Candidate.
      c = jnp.dot(jnp.concatenate([x, r * gru_state], axis=-1), w2) + b2
  
      new_gru_state = u * gru_state + (1 - u) * jnp.tanh(c)
      return new_gru_state, new_gru_state
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import activation_fns
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
label_desc: Use image node_id 2 for calling the GRUCell obj's underlying member method code with example usage: obj.forward(inputs) and returns: new_gru_state,  new_gru_state