jax

1aeede73 - Merge pull request #4742 from bchetioui:reduce_top_k_tests