Add cosmological model for the conversation from red-shifted mass to source-frame mass

[CharmaineWONG2]

The inferred parameters from PE are red-shifted. If we want the source-frame parameter, we may need to add cosmological models. For bilby, they use astropy.cosmology.FlatLambdaCDM but it is not written in jax. We may need one that is written in jax.
https://docs.astropy.org/en/stable/api/astropy.cosmology.FlatLambdaCDM.html#astropy.cosmology.FlatLambdaCDM

[ColmTalbot]

https://wcosmo.readthedocs.io/en/latest/ is a (mostly-stable) jax port of some astropy cosmologies