flax

d94b5a37 - fix seq_length in rnn, added test for rnn