Use jax.block_until_ready in benchmarks

[xela-95]

In this PR I added two little changes:

• Use different random data for each batch (created from different PRNG keys)
• Use jax.block_until_ready to measure the time of the full computation and not only the asynchronous dispatch (see https://jax.readthedocs.io/en/latest/async_dispatch.html#async-dispatch)

---

📚 Documentation preview 📚: https://jaxsim--344.org.readthedocs.build//344/

[flferretti]

Thanks a lot @xela-95, this was very much needed! LGTM

Before merging, would you mind updating the PR title to be more specific to the changes, like "Use jax.block_until_ready in benchmarks" or something related to avoid measuring the dispatch? Thank you

[xela-95]

Thanks a lot @xela-95, this was very much needed! LGTM

Before merging, would you mind updating the PR title to be more specific to the changes, like "Use jax.block_until_ready in benchmarks" or something related to avoid measuring the dispatch? Thank you

Done! Rebasing and merging then :)