code: def GRUCell(units):
    
    return GeneralGRUCell(
        candidate_transform=lambda: core.Dense(units=units),
        memory_transform=combinators.Identity,
        gate_nonlinearity=core.Sigmoid,
        candidate_nonlinearity=core.Tanh)
docs: 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.