jax

dda0aa17 - Update tests/experimental_rnn_test.py

Update tests/experimental_rnn_test.py Co-authored-by: JD <jahandad@gmail.com> (cherry picked from commit f005d96cc027a16edfb99b7b3764b26335a82b31)