Bump up version to 0.2.1 (#435)

* Bump up version to 0.2.1

* address comments

* use tuple instead of set

.travis.yml

@@ -14,6 +14,8 @@ dist: xenial
 
 install:
     - pip install -U pip
+    # Keep track of pyro-api master branch
+    - pip install https://github.com/pyro-ppl/pyro-api/archive/master.zip
     - pip install jaxlib
     - pip install jax
     - pip install .[examples,test]

examples/baseball.py

@@ -176,7 +176,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Baseball batting average using HMC")
     parser.add_argument("-n", "--num-samples", nargs="?", default=3000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1500, type=int)

examples/bnn.py

@@ -129,7 +129,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Bayesian neural network example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=2000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)

examples/covtype.py

@@ -56,7 +56,7 @@ def main(args):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument('-n', '--num-samples', default=100, type=int, help='number of samples')
     parser.add_argument('--num-steps', default=10, type=int, help='number of steps (for "HMC")')

examples/funnel.py

@@ -78,7 +78,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Non-centered reparameterization example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)

examples/gp.py

@@ -120,7 +120,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)

examples/hmm.py

@@ -158,7 +158,7 @@ def main(args):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description='Semi-supervised Hidden Markov Model')
     parser.add_argument('--num-categories', default=3, type=int)
     parser.add_argument('--num-words', default=10, type=int)

examples/minipyro.py

@@ -55,7 +55,7 @@ def body_fn(i, val):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Mini Pyro demo")
     parser.add_argument("-f", "--full-pyro", action="store_true", default=False)
     parser.add_argument("-n", "--num-steps", default=1001, type=int)

examples/neutra.py

@@ -148,7 +148,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="NeuTra HMC")
     parser.add_argument('-n', '--num-samples', nargs='?', default=10000, type=int)
     parser.add_argument('--num-warmup', nargs='?', default=1000, type=int)

examples/sparse_regression.py

@@ -312,7 +312,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Gaussian Process example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=500, type=int)

examples/stochastic_volatility.py

@@ -75,7 +75,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Stochastic Volatility Model")
     parser.add_argument('-n', '--num-samples', nargs='?', default=3000, type=int)
     parser.add_argument('--num-warmup', nargs='?', default=1500, type=int)

examples/ucbadmit.py

@@ -101,7 +101,7 @@ def main(args):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description='UCBadmit gender discrimination using HMC')
     parser.add_argument('-n', '--num-samples', nargs='?', default=2000, type=int)
     parser.add_argument('--num-warmup', nargs='?', default=500, type=int)

examples/vae.py

@@ -122,7 +122,7 @@ def reconstruct_img(epoch, rng_key):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.2.0')
+    assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument('-n', '--num-epochs', default=20, type=int, help='number of training epochs')
     parser.add_argument('-lr', '--learning-rate', default=1.0e-3, type=float, help='learning rate')

notebooks/source/bayesian_regression.ipynb

@@ -64,7 +64,7 @@
     "\n",
     "plt.style.use('bmh')\n",
     "\n",
-    "assert numpyro.__version__.startswith('0.2.0')"
+    "assert numpyro.__version__.startswith('0.2.1')"
    ]
   },
   {

notebooks/source/logistic_regression.ipynb

@@ -31,7 +31,7 @@
     "import numpyro.distributions as dist\n",
     "from numpyro.examples.datasets import COVTYPE, load_dataset\n",
     "from numpyro.infer import HMC, MCMC, NUTS\n",
-    "assert numpyro.__version__.startswith('0.2.0')\n",
+    "assert numpyro.__version__.startswith('0.2.1')\n",
     "\n",
     "# NB: replace gpu by cpu to run this notebook in cpu\n",
     "numpyro.set_platform(\"gpu\")"

notebooks/source/time_series_forecasting.ipynb

@@ -37,7 +37,7 @@
     "%matplotlib inline\n",
     "%config InlineBackend.figure_format = 'svg'\n",
     "\n",
-    "assert numpyro.__version__.startswith('0.2.0')"
+    "assert numpyro.__version__.startswith('0.2.1')"
    ]
   },
   {

numpyro/infer/mcmc.py

@@ -601,14 +601,14 @@ def _single_chain_mcmc(self, init, collect_fields=('z',), collect_warmup=False,
         states['z'] = vmap(constrain_fn)(states['z']) if len(tree_flatten(states['z'])[0]) > 0 else states['z']
         return states
 
-    def run(self, rng_key, *args, extra_fields=('diverging',), collect_warmup=False, init_params=None, **kwargs):
+    def run(self, rng_key, *args, extra_fields=(), collect_warmup=False, init_params=None, **kwargs):
         """
         Run the MCMC samplers and collect samples.
 
         :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
         :param args: Arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init` method.
             These are typically the arguments needed by the `model`.
-        :param extra_fields: Extra fields (aside from `z`) from :data:`numpyro.infer.mcmc.HMCState`
+        :param extra_fields: Extra fields (aside from `z`, `diverging`) from :data:`numpyro.infer.mcmc.HMCState`
             to collect during the MCMC run.
         :type extra_fields: tuple or list
         :param bool collect_warmup: Whether to collect samples from the warmup phase. Defaults
@@ -635,7 +635,7 @@ def run(self, rng_key, *args, extra_fields=('diverging',), collect_warmup=False,
                 raise ValueError('`init_params` must have the same leading dimension'
                                  ' as `num_chains`.')
         assert isinstance(extra_fields, (tuple, list))
-        collect_fields = ('z',) + tuple(extra_fields) if 'z' not in extra_fields else extra_fields
+        collect_fields = tuple(set(('z', 'diverging') + tuple(extra_fields)))
         if self.num_chains == 1:
             states_flat = self._single_chain_mcmc((rng_key, init_params), collect_fields, collect_warmup,
                                                   args, kwargs)

numpyro/version.py

@@ -1 +1 @@
-__version__ = '0.2.0'
+__version__ = '0.2.1'

setup.py

@@ -40,7 +40,7 @@
         'test': [
             'flake8',
             'pytest>=4.1',
-            'pyro-api@https://api.github.com/repos/pyro-ppl/pyro-api/tarball/master'
+            'pyro-api>=0.1.1'
         ],
         'dev': ['ipython'],
         'examples': ['matplotlib'],