Two parametrization of JAX benchmark test are taking hours

[ricardoV94]

Description

I thought that there was a problem in a recent PR I did, but I found others passed after HOURS 🗡️

Need to check if there was a recent regression introduced by us or JAX

[ricardoV94]

Two weeks ago the whole job was taking 16 minutes:

And last week jumped to ~3 hours (it's at 4 hours now):

[ricardoV94]

I suspect it is a JAX issue. I see widely different results across jax versions, although funnily enough 0.4.31 is the slow one on my machine and 0.4.34 a tad faster (reversed on the CI). 0.4.35 segfaults entirely, so there's that :)