jax

09a0ed6e - Merge pull request #4921 from bchetioui:fix_argminmax_gpu