Open AndreV84 opened 7 months ago
Could you add on a support for modified gravity model , please? We tried to accelerate https://miguelzuma.github.io/hi_class_public/ with jax GPU but it was not possible to get performance increase fro sampling https://discourse.pymc.io/t/hamiltonian-hmc-code-with-pymc-jax-gpu-sampler/13600/4
maybe you could share some directions or it can be implemented with jax cosmo somwhow?
Could you add on a support for modified gravity model , please? We tried to accelerate https://miguelzuma.github.io/hi_class_public/ with jax GPU but it was not possible to get performance increase fro sampling https://discourse.pymc.io/t/hamiltonian-hmc-code-with-pymc-jax-gpu-sampler/13600/4
maybe you could share some directions or it can be implemented with jax cosmo somwhow?