# NUTS sampling with GPU acceleration in PyMC4

**URL:** <https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791>\
**Category:** General\
**Tags:** gpu, hierarchical\
**Created:** [June 30, 2022, 5:26pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791 "2022-06-30T17:26:27Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![payamphysics](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/payamphysics/32/5218_2.png) [@payamphysics](https://discourse.pymc.io/u/payamphysics)\
**Post date:** [June 30, 2022, 5:26pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/1 "2022-06-30T17:26:27Z")

</div>

Hi there, I have set up a Hierarchical Bayes model for choice data (on AWS Sagemaker) and am able to use NUTS sampler in PyMC4 to take samples. Now I’m trying to run the sampling on GPU. My Sagemaker instance has GPU available. I tried to get it to work using a .aesara.rc file by setting device=cuda/cuda0/gpu but none of these work and only device=cpu works. I hear that GPU acceleration is supposed to be straightforward on PyMC4. Could someone please guide me in the right direction for that? More details of my model and sampling can be found here:

> [@Hierarchical Bayesian Choice modeling with PYMC4](https://discourse.pymc.io/t/hierarchical-bayesian-choice-modeling-with-pymc4/9733/7):
>
> For the record, here’s the code that solved the issue: basic\_model = pm.Model() with basic\_model: # Priors for unknown model parameters alphas = pm.Normal("alphas", mu=0, sigma=1, shape=num\_items) #sd\_distrib = pm.Exponential.dist(100.0) #sd\_distrib = pm.HalfNormal.dist(10) sd\_distrib = pm.Normal.dist(mu=1, sigma=10, shape=num\_items) #sd\_distrib = pm.ChiSquared.dist(nu=1, shape=num\_items) #sd\_distrib = pm.HalfCauchy.dist(beta=1, shape=num\_items) L\_sigma, co…

I can also provide more information if need be.  
Thanks in advance!

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 11, 2022, 2:28pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/2 "2022-07-11T14:28:56Z")

</div>

Have you followed the instructions on setting up your environment to use the GPU through the JAX backend?

> **[Set up JAX sampling with GPUs in PyMC v4 · pymc-devs/pymc Wiki](https://github.com/pymc-devs/pymc/wiki/Set-up-JAX-sampling-with-GPUs-in-PyMC-v4)**
>
> Probabilistic Programming in Python: Bayesian Modeling and Probabilistic Machine Learning with Aesara - Set up JAX sampling with GPUs in PyMC v4 · pymc-devs/pymc Wiki

You can check that JAX has found the GPU using:

```python
In [1]: import jax
In [2]: jax.default_backend()
Out[2]: 'gpu'
In [3]: jax.devices()
Out[3]: [GpuDevice(id=0, process_index=0)]

```

---

<div class="post-metadata">

**Author:** ![payamphysics](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/payamphysics/32/5218_2.png) [@payamphysics](https://discourse.pymc.io/u/payamphysics)\
**Post date:** [July 11, 2022, 5:58pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/3 "2022-07-11T17:58:50Z")

</div>

Yes, recently I followed those instructions and was able to confirm that gpu was being accessed, as you have suggested to check. Thank you!

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 11, 2022, 7:41pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/4 "2022-07-11T19:41:02Z")

</div>

Great, and you’re calling `pymc.sampling_jax.sample_numpyro_nuts()` to do your sampling? Sorry if these seem like basic questions, just trying to make sure we’re on the same page.

[https://www.pymc.io/projects/docs/en/stable/api/generated/pymc.sampling\_jax.sample\_numpyro\_nuts.html](https://www.pymc.io/projects/docs/en/stable/api/generated/pymc.sampling_jax.sample_numpyro_nuts.html)

---

<div class="post-metadata">

**Author:** ![payamphysics](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/payamphysics/32/5218_2.png) [@payamphysics](https://discourse.pymc.io/u/payamphysics)\
**Post date:** [July 11, 2022, 9:56pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/5 "2022-07-11T21:56:55Z")

</div>

No worries 🙂 Yes, that’s what I do.  
Thanks!

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 13, 2022, 8:24pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/6 "2022-07-13T20:24:48Z")

</div>

And how have you confirmed that the GPU isn’t being utilized?

---

<div class="post-metadata">

**Author:** ![payamphysics](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/payamphysics/32/5218_2.png) [@payamphysics](https://discourse.pymc.io/u/payamphysics)\
**Post date:** [July 13, 2022, 8:29pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/7 "2022-07-13T20:29:11Z")

</div>

I have realized that GPU “is” being utilized. Please see earlier post on this thread. Apparently, once JAX has GPU access, you’re done. PyMC4 utilizes GPU via JAX.

---

<div class="post-metadata">

**Author:** ![jhrcook](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/jhrcook/32/3255_2.png) [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Post date:** [July 13, 2022, 8:45pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/8 "2022-07-13T20:45:35Z")

</div>

Great! Glad you got it working. And yeah, when I was setting up to use the GPU, I was preparing for a fight with the computer, too. It was far easier than I expected.

---

<div class="post-metadata">

**Author:** ![payamphysics](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/payamphysics/32/5218_2.png) [@payamphysics](https://discourse.pymc.io/u/payamphysics)\
**Post date:** [July 14, 2022, 12:52am UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/9 "2022-07-14T00:52:45Z")

</div>

It was in fact a huge challenge for me.

---

<div class="post-metadata">

**Author:** ![Varun\_Gupta](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/varun_gupta/32/6345_2.png) [@Varun\_Gupta](https://discourse.pymc.io/u/Varun_Gupta)\
**Post date:** [July 6, 2023, 7:01pm UTC](https://discourse.pymc.io/t/nuts-sampling-with-gpu-acceleration-in-pymc4/9791/10 "2023-07-06T19:01:21Z")

</div>

Hi @payamphysics !

Hope you’re doing good !  
I’m trying to achieve what you did, but JAX wont detect GPU on AWS Sagemaker. Could you please tell me what instance/image did you use?