Repository: google/flax

Issue:
RNN called with seq_lengths gives a ValueError
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
Mac OS 11.7
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
jax==0.4.13
jaxlib==0.4.13
flax==0.7.0 (installed from github@main=881c4495a2423d5014ac0175dd490dfa2b6d5ab6)
- Python version:
Python 3.9.13

### Problem you have encountered:

Calling an RNN with a seq_lengths argument gives a ValueError

### What you expected to happen:

The RNN performs a forward pass

### Logs, error messages, etc:

```
...
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../venv/lib/python3.9/site-packages/flax/linen/recurrent.py", line 813, in __call__
    scan_output = scan(self.cell, carry, inputs)
  File ".../venv/lib/python3.9/site-packages/flax/core/axes_scan.py", line 139, in scan_fn
    _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)
  File ".../venv/lib/python3.9/site-packages/flax/core/axes_scan.py", line 118, in body_fn
    ys = jax.tree_util.tree_map(
ValueError: Expected tuple, got Traced<ShapedArray(float32[2,265]):JaxprTrace(level=1/0)>.
```
### Steps to reproduce:

```python
import jax
import jax.numpy as jnp
import flax.linen as nn

x = jnp.ones((2, 10, 6))
lstm = nn.RNN(nn.LSTMCell(265))
variables = lstm.init(jax.random.PRNGKey(0), x)
y = lstm.apply(variables, x, seq_lengths=jnp.array([5, 5]))
```


Discussion:
Ah, just realised that it works when `return_carry=True`

Resolving patch:
diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
--- a/flax/linen/recurrent.py
+++ b/flax/linen/recurrent.py
@@ -786,6 +786,8 @@ def __call__(
     else:
       carry = initial_carry
 
+    slice_carry = seq_lengths is not None and return_carry
+
     def scan_fn(
       cell: RNNCellBase, carry: Carry, x: Array
     ) -> Union[Tuple[Carry, Array], Tuple[Carry, Tuple[Carry, Array]]]:
@@ -794,7 +796,7 @@ def scan_fn(
       # so that we can select the last carry for each sequence later.
       # This uses more memory but is faster than using jnp.where at each
       # iteration. As a small optimization do this when we really need it.
-      if seq_lengths is not None and return_carry:
+      if slice_carry:
         return carry, (carry, y)
       else:
         return carry, y
@@ -802,7 +804,7 @@ def scan_fn(
     scan = transforms.scan(
       scan_fn,
       in_axes=time_axis,
-      out_axes=time_axis if seq_lengths is None else (0, time_axis),
+      out_axes=(0, time_axis) if slice_carry else time_axis,
       unroll=self.unroll,
       variable_axes=self.variable_axes,
       variable_broadcast=self.variable_broadcast,
@@ -815,7 +817,8 @@ def scan_fn(
     # Next we select the final carry. If a segmentation mask was provided and
     # return_carry is True we slice the carry history and select the last valid
     # carry for each sequence. Otherwise we just use the last carry.
-    if seq_lengths is not None and return_carry:
+    if slice_carry:
+      assert seq_lengths is not None
       _, (carries, outputs) = scan_output
       # seq_lengths[None] expands the shape of the mask to match the
       # number of dimensions of the carry.


Test patch:
diff --git a/tests/linen/linen_recurrent_test.py b/tests/linen/linen_recurrent_test.py
--- a/tests/linen/linen_recurrent_test.py
+++ b/tests/linen/linen_recurrent_test.py
@@ -369,6 +369,14 @@ def test_flip_sequence_time_major_more_feature_dims(self):
     np.testing.assert_allclose(flipped[:4, 0], x[:4, 0][::-1])
     np.testing.assert_allclose(flipped[:2, 1], x[:2, 1][::-1])
 
+  def test_basic_seq_lengths(self):
+
+    x = jnp.ones((2, 10, 6))
+    lstm = nn.RNN(nn.LSTMCell(265))
+    variables = lstm.init(jax.random.PRNGKey(0), x)
+    y = lstm.apply(variables, x, seq_lengths=jnp.array([5, 5]))
+
+
 class BidirectionalTest(absltest.TestCase):
 
   def test_bidirectional(self):