Apple M1 GPU in pymc4?

[fbarfi]

Was wondering if there is a chance to use Apple M1 GPU in pymc4 since Jax is used in pymc4 (though sampling_jax) and jax now recognizes Apple M1 GPU? PyTorch also recognizes now Apple M1 GPU, is it true for numpyro? and if that is true, can this be transferred to pymc4? thanks.

[ricardoV94]

PyMC should not have to do anything related to this specific backend. Is anything special needed on the JAX side for the GPU to be used? If not, it should work out of the box.

[fbarfi]

Thanks Ricardo. Whenever I use Jax directly it does detect the M1 GPU. However, the GPU is not detected when I use pm.sampling_jax. Is there a configuration that I need to set for that? Is it perhaps on aesara side?

[ricardoV94]

@fbarfi do you mind opening an issue on our discourse instead as it sounds like a user (not development) issue at the moment? Feel free to tag me there directly.

I'll close this discussion, but in the meantime I suggest you try to run a numpyro model directly and if it still doesn't catch the GPU, I would raise a discussion with them.

If it does, it may be that the jaxlib and jax versions you are using in one place may not match the ones in your pymc environment. Aesara should not be doing anything that prevents it.

[ricardoV94]

http://discourse.pymc.io/