jax

5b3e12a7 - Enabled RNN unit test: test_no_workspace_overflow for ROCm devices (#624)