Jax support

[aelkhour]

I understood from a quick discussion with @perrin-isir that JAX is becoming the standard for high-performance machine learning research, offering significant performance gains compared to pytorch and tensorflow.

I was wondering whether pinocchio can be interfaced as-is with JAX, or if some additional development work is needed.

/cc @duburcqa

[jcarpent]

We already support Just-In-Time paradigm for CPU, see the paper.
We will add support for GPUs and FPGAs.
It will be also supported in Python.

I don't think using JAX will lead to better performances in this aspect, mostly because the library has been tailored for ML applications, not properly supporting operations on Lie algebra. In addition, JAX focuses on Python usage, meaning it would be less easy to integrate than classic JIT frameworks.
More investigation on this aspect is of course needed.

External help is more than welcome.