Your JAX runs look slow – I’ve tried you code on an oldish CPU and runs significantly faster. A large overhead is being introduced for some reason. I would first try calling the JAX jit method on the jacobian and saving the resulting function to a variable which is then called.
With JIT and jacrev on my CPU (E5-2680 v3 @ 2.50GHz) I get time of 0.393. NB you should call block_until_ready() on last output to ensure asynchronous execution does not cause misleading results.