Bump to 0.6.0 (#959)

* sketch the plan

* add a running implementation (not working yet

* remove unnecessary changes

* sequential update

* temp save

* add a working implementation

* add 24d example

* fix lint

* test for the order

* merge master

* fix bug at reset momentum

* add various mh functions

* add various discrete gibbs function method

* change stay_prob to modified to avoid confusing users

* expose more information for mixed hmc

* sketch an implementation

* temp save

* temp save

* finish the implementation

* keep kinetic energy

* add temperature experiment

* add dual averaging

* add various debug statements

* fix bugs

* clean up and separating out clock adapter; but target distribution is wrong due to a bug somewhere

* clean up

* add comments and an example

* make sure forward mode work

* add docs for new HMC fields

* add tests for mixedhmc

* fix step_size bug

* use modified=False

* tests pass with the fix

* skip print summary

* adjust trajectory length

* port update_version script from Pyro

* pin jax/jaxlib versions

* run isort

* fix some issues during collection notes

* use result_type instead of canonicalize_dtype

* fix lint

* change get_dtype  to jnp.result_type

* add print summary

* fix compiling issue for mcmc

* also try to avoid compiling issue in other samplers

* also fix compiling issue in barkermh

* convert init params types to strong types

* address comments

* fix wrong docs

* run isort

Author: fehiepsi

README.md

@@ -182,7 +182,7 @@ Pyro users will note that the API for model specification and inference is large
 
 ## Installation
 
-> **Limited Windows Support:** Note that NumPyro is untested on Windows, and might require building jaxlib from source. See this [JAX issue](https://github.com/google/jax/issues/438) for more details. Alternatively, you can install [Windows Subsystem for Linux](https://docs.microsoft.com/en-us/windows/wsl/) and use NumPyro on it as on a Linux system. See also [CUDA on Windows Subsystem for Linux](https://developer.nvidia.com/cuda/wsl) if you want to use GPUs on Windows.
+> **Limited Windows Support:** Note that NumPyro is untested on Windows, and might require building jaxlib from source. See this [JAX issue](https://github.com/google/jax/issues/438) for more details. Alternatively, you can install [Windows Subsystem for Linux](https://docs.microsoft.com/en-us/windows/wsl/) and use NumPyro on it as on a Linux system. See also [CUDA on Windows Subsystem for Linux](https://developer.nvidia.com/cuda/wsl) and [this forum post](https://forum.pyro.ai/t/numpyro-with-gpu-works-on-windows/2690) if you want to use GPUs on Windows.
 
 To install NumPyro with a CPU version of JAX, you can use pip:
 

docs/source/distributions.rst

@@ -574,6 +574,14 @@ real_vector
 -----------
 .. autodata:: numpyro.distributions.constraints.real_vector
 
+softplus_positive
+-----------------
+.. autodata:: numpyro.distributions.constraints.softplus_positive
+
+softplus_lower_cholesky
+-----------------------
+.. autodata:: numpyro.distributions.constraints.softplus_lower_cholesky
+
 simplex
 -------
 .. autodata:: numpyro.distributions.constraints.simplex

docs/source/mcmc.rst

@@ -70,6 +70,8 @@ MCMC Kernels
 
 .. autofunction:: numpyro.infer.hmc.hmc.sample_kernel
 
+.. autofunction:: numpyro.infer.hmc_gibbs.taylor_proxy
+
 .. autodata:: numpyro.infer.barker.BarkerMHState
 
 .. autodata:: numpyro.infer.hmc.HMCState

examples/annotation.py

@@ -266,7 +266,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith("0.5.0")
+    assert numpyro.__version__.startswith('0.6.0')
     parser = argparse.ArgumentParser(description="Bayesian Models of Annotation")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)

examples/baseball.py

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

examples/bnn.py

@@ -138,7 +138,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.5.0')
+    assert numpyro.__version__.startswith('0.6.0')
     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

@@ -139,7 +139,7 @@ def main(args):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.5.0')
+    assert numpyro.__version__.startswith('0.6.0')
     parser = argparse.ArgumentParser(description="parse args")
     parser.add_argument('-n', '--num-samples', default=1000, type=int, help='number of samples')
     parser.add_argument('--num-warmup', default=1000, type=int, help='number of warmup steps')

examples/funnel.py

@@ -87,7 +87,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.5.0')
+    assert numpyro.__version__.startswith('0.6.0')
     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

@@ -142,7 +142,7 @@ def main(args):
 
 
 if __name__ == "__main__":
-    assert numpyro.__version__.startswith('0.5.0')
+    assert numpyro.__version__.startswith('0.6.0')
     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

@@ -191,7 +191,7 @@ def main(args):
 
 
 if __name__ == '__main__':
-    assert numpyro.__version__.startswith('0.5.0')
+    assert numpyro.__version__.startswith('0.6.0')
     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)