Add num_chains arg to examples/.../sir.py

pyro-ppl · Jun 15, 2020 · 30d7a05 · 30d7a05
1 parent 45f2e7b
commit 30d7a05
Show file tree

Hide file tree

Showing 3 changed files with 386 additions and 0 deletions.
diff --git a/examples/contrib/epidemiology/sir.py b/examples/contrib/epidemiology/sir.py
@@ -91,6 +91,7 @@ def hook_fn(kernel, *unused):
                      heuristic_ess_threshold=args.ess_threshold,
                      warmup_steps=args.warmup_steps,
                      num_samples=args.num_samples,
+                     num_chains=args.num_chains,
                      max_tree_depth=args.max_tree_depth,
                      arrowhead_mass=args.arrowhead_mass,
                      num_quant_bins=args.num_bins,
@@ -293,6 +294,7 @@ def main(args):
     parser.add_argument("-np", "--num-particles", default=1024, type=int)
     parser.add_argument("-ess", "--ess-threshold", default=0.5, type=float)
     parser.add_argument("-w", "--warmup-steps", type=int)
+    parser.add_argument("-c", "--num-chains", default=1, type=int)
     parser.add_argument("-t", "--max-tree-depth", default=5, type=int)
     parser.add_argument("-a", "--arrowhead-mass", action="store_true")
     parser.add_argument("-r", "--rng-seed", default=0, type=int)

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -34,6 +34,7 @@
     'contrib/autoname/tree_data.py --num-epochs=1',
     'contrib/cevae/synthetic.py --num-epochs=1',
     'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2',
+    'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -c=2',
     'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -e=2',
     'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -k=1',
     'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -e=2 -k=1',