From 9367ba32716b9e6334d2a149f686f9787a682be4 Mon Sep 17 00:00:00 2001
From: damonbayer <xum8@cdc.gov>
Date: Thu, 18 Jul 2024 13:46:45 -0400
Subject: [PATCH] set host device count

---
 docs/source/tutorials/basic_renewal_model.qmd    |  6 +++---
 .../tutorials/hospital_admissions_model.qmd      | 16 ++++++++++------
 docs/source/tutorials/pyrenew_demo.qmd           |  5 ++++-
 3 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/docs/source/tutorials/basic_renewal_model.qmd b/docs/source/tutorials/basic_renewal_model.qmd
index 2e324a08..ace4a3f7 100644
--- a/docs/source/tutorials/basic_renewal_model.qmd
+++ b/docs/source/tutorials/basic_renewal_model.qmd
@@ -27,6 +27,8 @@ from pyrenew.deterministic import DeterministicPMF
 from pyrenew.model import RtInfectionsRenewalModel
 from pyrenew.metaclass import DistributionalRV
 import pyrenew.transformation as t
+
+npro.set_host_device_count(2)
 ```
 
 ## Architecture of `RtInfectionsRenewalModel`
@@ -214,9 +216,7 @@ model1.run(
     num_samples=1000,
     data_observed_infections=sim_data.observed_infections,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(
-        progress_bar=False, num_chains=2, chain_method="sequential"
-    ),
+    mcmc_args=dict(progress_bar=False, num_chains=2),
 )
 ```
 
diff --git a/docs/source/tutorials/hospital_admissions_model.qmd b/docs/source/tutorials/hospital_admissions_model.qmd
index f550ea98..d1e00c8e 100644
--- a/docs/source/tutorials/hospital_admissions_model.qmd
+++ b/docs/source/tutorials/hospital_admissions_model.qmd
@@ -4,6 +4,14 @@ format: gfm
 engine: jupyter
 ---
 
+```{python}
+# | label: numpyro setup
+# | echo: false
+import numpyro as npro
+
+npro.set_host_device_count(2)
+```
+
 This document illustrates how a hospital admissions-only model can be fitted using data from the Pyrenew package, particularly the wastewater dataset. The CFA wastewater team created this dataset, which contains simulated data.
 
 ## Model definition
@@ -205,7 +213,6 @@ Let's simulate to check if the model is working:
 
 ```{python}
 # | label: simulation
-import numpyro as npro
 import numpy as np
 
 timeframe = 120
@@ -245,15 +252,12 @@ We can fit the model to the data. We will use the `run` method of the model obje
 # | label: model-fit
 import jax
 
-npro.set_host_device_count(jax.local_device_count())
 hosp_model.run(
     num_samples=1000,
     num_warmup=1000,
     data_observed_hosp_admissions=daily_hosp_admits,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(
-        progress_bar=False, num_chains=2, chain_method="sequential"
-    ),
+    mcmc_args=dict(progress_bar=False, num_chains=2),
 )
 ```
 
@@ -535,7 +539,7 @@ hosp_model_weekday.run(
     num_warmup=2000,
     data_observed_hosp_admissions=daily_hosp_admits,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(progress_bar=False),
+    mcmc_args=dict(progress_bar=False, num_chains=2),
     padding=pad_size,
 )
 ```
diff --git a/docs/source/tutorials/pyrenew_demo.qmd b/docs/source/tutorials/pyrenew_demo.qmd
index 858b41a6..d7444d70 100644
--- a/docs/source/tutorials/pyrenew_demo.qmd
+++ b/docs/source/tutorials/pyrenew_demo.qmd
@@ -29,8 +29,11 @@ import matplotlib.pyplot as plt
 import jax
 import jax.numpy as jnp
 import numpy as np
+import numpyro
 from numpyro.handlers import seed
 import numpyro.distributions as dist
+
+numpyro.set_host_device_count(2)
 ```
 
 ```{python}
@@ -180,7 +183,7 @@ hospmodel.run(
     num_samples=1000,
     data_observed_hosp_admissions=x.observed_hosp_admissions,
     rng_key=jax.random.PRNGKey(54),
-    mcmc_args=dict(progress_bar=False),
+    mcmc_args=dict(progress_bar=False, num_chains=2),
 )
 ```