Implement metric scaling

[AdrienCorenflos]

Thank you for opening a PR!

This is a change that relates to the PR #731.
The idea is to allow metrics to scale a vector (possibly momentum, possibly otherwise) by the cholesky of the mass matrix or its inverse.

DRAFT PR so far, taking comments on interface and tentative implementation. Nothing is tested yet.

A comment: for efficiency reasons (in the context of Riemannian metrics), I have allowed the scale function to take a sequence of elements to rescale, so as to avoid recomputing the mass matrix all the time. There may be a cleaner way to do this but I can't see it.

[AdrienCorenflos]

@ismael-mendoza this is what you need right?

[junpenglao]

A comment: for efficiency reasons (in the context of Riemannian metrics), I have allowed the scale function to take a sequence of elements to rescale, so as to avoid recomputing the mass matrix all the time. There may be a cleaner way to do this but I can't see it.

I see... I dont love it but I also dont hate it LOL

[ismael-mendoza]

Thanks @AdrienCorenflos I think this will work great, happy to incoporate this into #731 once it's merged!

[AdrienCorenflos]

I am a bit unhappy with one design choice: for the function scale to mean the same thing for all metrics I had to treat inputs differently: this comes from the fact that euclidean metrics take an inverse mass_matrix, while Riemannian take the mass matrix... I'm not sure what the rationale was behind this, but we will need to rethink it down the line. Maybe Riemannian should also take the inverse mass matrix? Either way, not a problem for this PR.

[junpenglao]

A likely over-engineered solution here is to create a new covariance matrix class that has both original and inverse form, kind of like what TFP does for Gaussian Process https://github.com/tensorflow/probability/blob/main/tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel.py

[AdrienCorenflos]

I changed the api a bit -> you now only give scale(position, vector, inv) rather than something complicated. This is because I remembered we are using JAX and I shouldn't reinvent the wheel: if people want several scalings they should use vmap and that's it. This will still only call mass_matrix (and cholesky etc) only once.

[AdrienCorenflos]

@junpenglao any comments or go?

[junpenglao]

Let me take another look tomorrow

[AdrienCorenflos]

That should be good now @junpenglao
Sorry about the hold up @ismael-mendoza but I think this is for a good reason :D

[junpenglao]

Great stuff, thank you!