From 25b71e6c7cd1a84db90ff2f314667d33cb3f385c Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Sun, 14 Jun 2020 10:25:32 -0700 Subject: [PATCH] Enable jit by default in examples --- examples/contrib/epidemiology/regional.py | 3 +- examples/contrib/epidemiology/sir.py | 3 +- tests/test_examples.py | 40 +++++++++++------------ 3 files changed, 24 insertions(+), 22 deletions(-) diff --git a/examples/contrib/epidemiology/regional.py b/examples/contrib/epidemiology/regional.py index 0328a57ced..1f2b39210a 100644 --- a/examples/contrib/epidemiology/regional.py +++ b/examples/contrib/epidemiology/regional.py @@ -150,7 +150,8 @@ def main(args): parser.add_argument("--single", action="store_false", dest="double") parser.add_argument("--rng-seed", default=0, type=int) parser.add_argument("--cuda", action="store_true") - parser.add_argument("--jit", action="store_true") + parser.add_argument("--jit", action="store_true", default=True) + parser.add_argument("--nojit", action="store_true", dest="jit") parser.add_argument("--verbose", action="store_true") parser.add_argument("--plot", action="store_true") args = parser.parse_args() diff --git a/examples/contrib/epidemiology/sir.py b/examples/contrib/epidemiology/sir.py index 6fb5bd0bb0..25417f4c94 100644 --- a/examples/contrib/epidemiology/sir.py +++ b/examples/contrib/epidemiology/sir.py @@ -300,7 +300,8 @@ def main(args): parser.add_argument("--double", action="store_true", default=True) parser.add_argument("--single", action="store_false", dest="double") parser.add_argument("--cuda", action="store_true") - parser.add_argument("--jit", action="store_true") + parser.add_argument("--jit", action="store_true", default=True) + parser.add_argument("--nojit", action="store_true", dest="jit") parser.add_argument("--verbose", action="store_true") parser.add_argument("--plot", action="store_true") args = parser.parse_args() diff --git a/tests/test_examples.py b/tests/test_examples.py index 838ae86d56..459266a43e 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -33,19 +33,19 @@ 'contrib/autoname/mixture.py --num-epochs=1', 'contrib/autoname/tree_data.py --num-epochs=1', 'contrib/cevae/synthetic.py --num-epochs=1', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -e=2', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -k=1', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -e=2 -k=1', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --haar', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -nb=4', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -hfm=3', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -a', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -o=0.2', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --haar', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 -hfm=3', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 -nb=4', + '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 -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', + 'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --haar', + 'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -nb=4', + 'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -hfm=3', + 'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -a', + 'contrib/epidemiology/sir.py --nojit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -o=0.2', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --haar', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 -hfm=3', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 -nb=4', 'contrib/forecast/bart.py --num-steps=2 --stride=99999', 'contrib/gp/sv-dkl.py --epochs=1 --num-inducing=4 --batch-size=1000', 'contrib/gp/sv-dkl.py --binary --epochs=1 --num-inducing=4 --batch-size=1000', @@ -108,11 +108,11 @@ 'air/main.py --num-steps=1 --cuda', 'baseball.py --num-samples=200 --warmup-steps=100 --num-chains=2 --cuda', 'contrib/cevae/synthetic.py --num-epochs=1 --cuda', - 'contrib/epidemiology/sir.py -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --cuda', - 'contrib/epidemiology/sir.py -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -nb=16 --cuda', - 'contrib/epidemiology/sir.py -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --haar --cuda', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --cuda', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --haar --cuda', + 'contrib/epidemiology/sir.py --nojit -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --cuda', + 'contrib/epidemiology/sir.py --nojit -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 -nb=16 --cuda', + 'contrib/epidemiology/sir.py --nojit -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --haar --cuda', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --cuda', + 'contrib/epidemiology/regional.py --nojit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --haar --cuda', 'contrib/gp/sv-dkl.py --epochs=1 --num-inducing=4 --cuda', 'lkj.py --n=50 --num-chains=1 --warmup-steps=100 --num-samples=200 --cuda', 'dmm/dmm.py --num-epochs=1 --cuda', @@ -157,8 +157,8 @@ def xfail_jit(*args): 'baseball.py --num-samples=200 --warmup-steps=100 --jit', 'contrib/autoname/mixture.py --num-epochs=1 --jit', 'contrib/cevae/synthetic.py --num-epochs=1 --jit', - 'contrib/epidemiology/sir.py -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2 --jit', - 'contrib/epidemiology/regional.py -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2 --jit', + 'contrib/epidemiology/sir.py --jit -np=128 -t=2 -w=2 -n=4 -d=20 -p=1000 -f 2', + 'contrib/epidemiology/regional.py --jit -t=2 -w=2 -n=4 -r=3 -d=20 -p=1000 -f 2', xfail_jit('lkj.py --n=50 --num-chains=1 --warmup-steps=100 --num-samples=200 --jit'), xfail_jit('contrib/gp/sv-dkl.py --epochs=1 --num-inducing=4 --jit'), xfail_jit('dmm/dmm.py --num-epochs=1 --jit'),