Make all fields of simulator.JaxSim static excluding SimulatorData

[diegoferigo]

With @flferretti, we noticed the occurrence of trace leaks in a jitted context of a parallel simulation on GPUs. It turns out that the leak was caused by the following line:

jaxsim/src/jaxsim/simulation/simulator.py

Line 359 in 6397508

dt_ns = jnp.array(self.step_size_ns * self.steps_per_run, dtype=jnp.uint64)

For some reason, jax detects a leak in this operation.

Converting JaxSim.step_size_ns to a static attribute was the most straightforward solution. The pay to price, this time, is losing the possibility to perform a parallel simulation whose models are integrated using different step sizes, that is regardless an edge case not so common. If needed, we can try to fix it in the future, what matters now is to simulate a large batch of models all together on hardware accelerators.

After this change, all attributes of JaxSim are marked as Static, with the exception of JaxSim.data that now contains the only mutable information — the real simulator state. It makes clear what can be changed during runtime —dynamically, w/o triggering recompilations— and what not.