Effect of argument seq_lengths in RNN module

[Kajiih]

Hello dear Flax community,
I am trying to understand the effect of passing seq_lengths argument when calling a RNN module when reverse is True compared to when it is False. I found it a little confusing in the documentation, so I would need some clarifications.

Here's what I understood:

• If have have a sequence of length T padded to in an array of length N.
  The output of the RNN (without reverse) will be exactly the same whether I use seq_lengths=T or not, is it correct? Only the carry will be different and in both cases, the valid data will be in the first T indexes of the output, the last one in a temporal sense, being at index T-1, right?

• However if I set reverse=True, the outputs will be different, and in particular, the only correct one is the one where I use seq_lengths=T (the other would use wrong data). And in that case, the valid data will be in the last T indexes of the output, i.e. N-T-1 to N-1 (and the last temporal one a index N-1). Also in the case I set keep_order=True, the output will be reversed and the valid data will again be in the first T indexes, but this time, the last (in a temporal sense) one will be in index 0.

If I understood everything correctly, in practice, setting seq_lengths is in fact useless unless we also use reverse or the final carry.

Sorry if it is rather confusing too, but these indexes stories are always a mess for me to manage... (also it's the first time I'm playing with RNNs so, I'm not really familiar with these)
Thanks in advance and thanks a lot for your great work and passion!

[cgarciae]

Hey @Kajiih, sorry for the late reply.

And in that case, the valid data will be in the last T indexes of the output, i.e. N-T-1 to N-1 (and the last temporal one a index N-1).

This is incorrect. The sequence will be reverse but padding will be kept at the end. For example if you have the sequence:

[ a, b, c, 0, 0 ]

0s meaning padding, it will be read as:

[ c, b, a, 0, 0 ]

I'm wondering if we should update our docstrings to reflect this?

[Kajiih]

Thank you @cgarciae for the clarification! So with reverse, valid data are in indexes 0 to T-1 too, I was probably confused because i was not sure of what exactly sequence was refering to in "If seq_lengths is passed, padding will always remain at the end of the sequence."

[ a, b, c, 0, 0 ]

0s meaning padding, it will be read as:
[ c, b, a, 0, 0 ]

I think a simple example like this would be enough to make it very clear. It would probably be beneficial, since it is somewhat hard to check this behavior directly as nothing is zeroed out.