Fixes distance computation to avoid unnecessary retracing

[EiffL]

This small fix leads to huge improvement in eager execution time, by compiling the ODE solving functions needed for distance computations. Also improved the speed of grothw ODE by reducing the number of necessary steps.

[EiffL]

This PR should close #70

[EiffL]

That's a good question ^^' Because I said before we shouldn't ^^'

So, in general we don't want to use tf.function except for the most external function, BUT in the case of the ODE solver, if the function that contains it is not compiled with a tf.function, we get these warnings about retracing, and it's very inefficient to use in eager mode (remember, that's why we had wrapping functions around the distance computation when we wanted to use them in the notebook).

So, in this particular case, we want to encapsulate the ODE solver in a tf.function, to make it efficient. But tf.function require very specific inputs, only tensors, so for instance a tf.function can't accept a dictionary or a cosmology object in principle. This makes them not very user friendly, and that's why we don't write everything as tf.functions.

So, what I did here, is that I created a "hidden" function, that the user will not have access to, _distance_computation_func which is compiled, and then the user friendly function rad_comoving_distance internally call that function.

Does that make sense?

[dlanzieri]

@EiffL Yes, you were very clear. Thank you for your time!