[Feature request]Unify the order of input arguments and returns ​​of RNN

Hi~
Recently, when I was using flax, I found that the order of the input arguments of RNN and RNNCell in flax is not uniform. This is very inconvenient when using it. I often need to think about which argument I should write first.

RNN call first accepts `inputs` arguments, then `carry` and other arguments, while RNNCell accepts `carry` first and then `inputs`. Similarly, the order of the return value of RNN call is inconsistent with the order of input arguments, which does not seem to be intuitive.

I know that the `carry` argument can be ignored in RNN call, so it is placed after the `inputs` argument, but this is not the case in RNNCell. Both `inputs` and `carry` are necessary argument in RNNCell, so can you adjust the input arguments in RNNCell? Make the order is consistent with the order in the RNN call, and the return values order of RNN and RNNCell is modified to be consistent with the input order. I think this would make flax simpler and easier to use, and as far as I know, RNNCell and rnn_unroll in haiku are implemented like this, and RNN and RNNCell in pytorch also implemented like this.

**The impemented in flax now:**
```python
class GRUCell(RNNCellBase):
...

@compact
def __call__(self, carry, inputs):
"""Gated recurrent unit (GRU) cell.

Args:
carry: the hidden state of the GRU cell,
initialized using ``GRUCell.initialize_carry``.
inputs: an ndarray with the input for the current time step.
All dimensions except the final are considered batch dimensions.

Returns:
A tuple with the new carry and the output.
"""
````
```python
class LSTMCell(RNNCellBase):
...

@compact
def __call__(self, carry, inputs):
r"""A long short-term memory (LSTM) cell.

Args:
carry: the hidden state of the LSTM cell,
initialized using ``LSTMCell.initialize_carry``.
inputs: an ndarray with the input for the current time step.
All dimensions except the final are considered batch dimensions.

Returns:
A tuple with the new carry and the output.
"""
```
```python
class RNN(Module):
...

def __call__(
self,
inputs: jax.Array,
*,
initial_carry: Optional[Carry] = None,
init_key: Optional[PRNGKey] = None,
seq_lengths: Optional[Array] = None,
return_carry: Optional[bool] = None,
time_major: Optional[bool] = None,
reverse: Optional[bool] = None,
keep_order: Optional[bool] = None,
) -> Union[Output, Tuple[Carry, Output]]:
"""
Applies the RNN to the inputs.

...

Returns:
if ``return_carry=False`` (default) only the output sequence is returned,
else it will return a tuple of the final carry and the output sequence.
"""
```
**which I want:**
```python
class GRUCell(RNNCellBase):
...

@compact
def __call__(self, inputs, carry):
"""Gated recurrent unit (GRU) cell.

Args:
inputs: an ndarray with the input for the current time step.
All dimensions except the final are considered batch dimensions.
carry: the hidden state of the GRU cell,
initialized using ``GRUCell.initialize_carry``.

Returns:
A tuple with the output and the new carry.
"""
```
```python
class LSTMCell(RNNCellBase):
...

@compact
def __call__(self, inputs, carry):
r"""A long short-term memory (LSTM) cell.

Args:
inputs: an ndarray with the input for the current time step.
All dimensions except the final are considered batch dimensions.
carry: the hidden state of the LSTM cell,
initialized using ``LSTMCell.initialize_carry``.

Returns:
A tuple with the output and the new carry.
"""
```
```python
class RNN(Module):
...

def __call__(
self,
inputs: jax.Array,
*,
initial_carry: Optional[Carry] = None,
init_key: Optional[PRNGKey] = None,
seq_lengths: Optional[Array] = None,
return_carry: Optional[bool] = None,
time_major: Optional[bool] = None,
reverse: Optional[bool] = None,
keep_order: Optional[bool] = None,
) -> Union[Output, Tuple[Carry, Output]]:
"""
Applies the RNN to the inputs.

...

Returns:
if ``return_carry=False`` (default) only the output sequence is returned,
else it will return a tuple of the output sequence and the final carry.
"""
```