Hello, I want to know if there is any new release of pathfinder running in GPU, because with pymc-extras still run on cpu.
Thanks.
You can set compile_kwargs={'mode':'JAX'} to put the log_dlogp function on GPU. The actual algorithm will be run on CPU, so you will pay the cost of going back and forth. blackjax offers are pure jax implementation of the end-to-end pathfinder algorithm, so you can bring a jax-compiled logp_dlogp function over there and get the whole thing on GPU. Full disclosure, I have never tried the blackjax route so I don’t know the API. I know @junpenglao has been upgrading it recently to include multi-path, so that’s cool.
1 Like
thank you for your response.