jax

4015de85 - Update tests/experimental_rnn_test.py

Update tests/experimental_rnn_test.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>