[linen] Undocumented RNN errors

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04 WSL
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: flax = 0.10.2 jax=0.5.0, jaxlib=0.5.0
- Python version: 12
- GPU/TPU model and memory: 1x Tesla V100, 32.0 GB
- CUDA version (if applicable):

### Problem you have encountered:
Hello Flax,
I'm stuck debugging linen.Bidirectional because it does not throw errors.
Logic tells me to do something like this:
```python
def setup(self):

self.rnn_var = nn.recurrent.Bidirectional(
nn.RNN(
nn.recurrent.OptimizedLSTMCell(self.rnn_var_lstm_cell_size),
return_carry=True,
),
nn.RNN(
nn.recurrent.OptimizedLSTMCell(self.rnn_var_lstm_cell_size),
return_carry=True,
),
return_carry=True,
) # num_layers=10,

self.initial_rnn_carry = (
( # note: don't forget to add ibatch size in the future.
nn.initializers.zeros(
jax.random.PRNGKey(0),
(self.rnn_var_lstm_cell_size,),
), # c (cell state)
nn.initializers.zeros(
jax.random.PRNGKey(0),
(self.rnn_var_lstm_cell_size,),
),
), # h (hidden state)
(
nn.initializers.zeros(
jax.random.PRNGKey(0),
(self.rnn_var_lstm_cell_size,),
), # c (cell state) for backward
nn.initializers.zeros(
jax.random.PRNGKey(0),
(self.rnn_var_lstm_cell_size,),
),
), # h (hidden state) for backward
)
# ...

def __call__(self, x: dict[str, Array]) -> Array:
import pdb

pdb.set_trace()
rnn_var_carry, _processed_vars = self.rnn_var_rnn(
nodes_gather_x, initial_carry=self.initial_rnn_carry, return_carry=True
)
#failure here^
```
And the problem is, I don't get understandable error messages, at all. For example this:
```
File "/Problemologist-flax/dreamerv3_flax/encoder.py", line 159, in __call__
rnn_var_rnn_carry, _processed_rnn_var = self.rnn_var_rnn(
^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/linen/recurrent.py", line 1315, in __call__
carry_forward, outputs_forward = self.forward_rnn(
^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/linen/recurrent.py", line 1135, in __call__
scan_output = scan(self.cell, carry, inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/core/axes_scan.py", line 152, in scan_fn
_, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/core/axes_scan.py", line 124, in body_fn
broadcast_out, c, ys = fn(broadcast_in, c, *xs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/linen/recurrent.py", line 1114, in scan_fn
carry, y = cell(carry, x)
^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/linen/recurrent.py", line 336, in __call__
dense_params_i[component] = DenseParams(
^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/flax/linen/recurrent.py", line 214, in __call__
(inputs.shape[-1], self.features),
~~~~~~~~~~~~^^^^ #