Builds a traditional GRU cell with dense internal transformations.

  Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555


  Args:
    units: Number of hidden units.

  Returns:
    A Stax model representing a traditional GRU RNN cell.

def GRUCell(units):
  """Builds a traditional GRU cell with dense internal transformations.

  Gated Recurrent Unit paper: https://arxiv.org/abs/1412.3555


  Args:
    units: Number of hidden units.

  Returns:
    A Stax model representing a traditional GRU RNN cell.
  """
  return GeneralGRUCell(
      candidate_transform=lambda: core.Dense(units=units),
      memory_transform=combinators.Identity,
      gate_nonlinearity=core.Sigmoid,
      candidate_nonlinearity=core.Tanh)