Speed up various functions avoiding using jax.vmap on indexes

[flferretti]

This PR removes the usage of jax.vmap on functions that can directly extract from arrays using the numpy-like indexing, e.g. using an index of an array for extracting and then stacking the resulting extracted components. This speeds up the compilation and the runtime performance of the concerned methods.

A couple benchmarks on stickbot model with 23 joints:

In [61]: %timeit js.contact.transforms(model=model, data=data) # previous implementation
11.9 ms ± 176 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [66]: %timeit js.contact.transforms(model=model, data=data) # current implementation
1.63 ms ± 45.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [98]: %timeit -r 1 -n 1 js.model.total_mass(model) # previous implementation
5.98 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [100]: %timeit -r 1 -n 1 js.model.total_mass(model) # current implementation
1.61 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

Apart from having a noticeable impact on the single methods, this should speed up all the functionalities that depend on the concerned functions.

---

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

[diegoferigo]

Thanks @flferretti, all these tiny performance improvements are more than welcome! All good, except the change in the usage of jaxsim.api.link.mass() that puzzles me a bit. I suggested a low-level workaround.