Maintain function that converts to torch?

[mjo22]

This could add a lot of value for people. See this gist: https://gist.github.com/mattjj/e8b51074fed081d765d2f3ff90edf0e9

It could simply go in cryojax.jax_util. The only thing is we don't want torch as a dependency; we would have to make it optional.