Add docs for handlers and distributions (#179)

* Add docs for handlers and distributions

* fix bugs; rebase

* fix doctest

* fix tests

* fix remaining bugs

* fix example

* address comments

* raise value error if algo isn't valid

* fix doc

* better organization

* fix arg in example

docs/source/diagnostics.rst

@@ -1,8 +1,27 @@
-Diagnostics
-===========
-
-.. automodule:: numpyro.diagnostics
-    :members:
-    :undoc-members:
-    :show-inheritance:
-    :member-order: bysource
+Autocorrelation
+---------------
+.. autofunction:: numpyro.diagnostics.autocorrelation
+
+Autocovariance
+--------------
+.. autofunction:: numpyro.diagnostics.autocovariance
+
+Effective Sample Size
+---------------------
+.. autofunction:: numpyro.diagnostics.effective_sample_size
+
+Gelman Rubin
+------------
+.. autofunction:: numpyro.diagnostics.gelman_rubin
+
+Split Gelman Rubin
+------------------
+.. autofunction:: numpyro.diagnostics.split_gelman_rubin
+
+HPDI
+----
+.. autofunction:: numpyro.diagnostics.hpdi
+
+Summary
+-------
+.. autofunction:: numpyro.diagnostics.summary

docs/source/distributions.rst

@@ -1,8 +1,250 @@
-Distributions
-=============
+Base Distribution
+=================
 
-.. automodule:: numpyro.distributions
+Distribution
+------------
+.. autoclass:: numpyro.distributions.Distribution
     :members:
     :undoc-members:
     :show-inheritance:
     :member-order: bysource
+
+TransformedDistribution
+-----------------------
+.. autoclass:: numpyro.distributions.TransformedDistribution
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+
+Continuous Distributions
+========================
+
+Beta
+----
+.. autoclass:: numpyro.distributions.continuous.Beta
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Cauchy
+------
+.. autoclass:: numpyro.distributions.continuous.Cauchy
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Chi2
+----
+.. autoclass:: numpyro.distributions.continuous.Chi2
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Dirichlet
+---------
+.. autoclass:: numpyro.distributions.continuous.Dirichlet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Exponential
+-----------
+.. autoclass:: numpyro.distributions.continuous.Exponential
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Gamma
+-----
+.. autoclass:: numpyro.distributions.continuous.Gamma
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+GaussianRandomWalk
+------------------
+.. autoclass:: numpyro.distributions.continuous.GaussianRandomWalk
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+HalfCauchy
+----------
+.. autoclass:: numpyro.distributions.continuous.HalfCauchy
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+HalfNormal
+----------
+.. autoclass:: numpyro.distributions.continuous.HalfNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+LKJCholesky
+-----------
+.. autoclass:: numpyro.distributions.continuous.LKJCholesky
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+LogNormal
+---------
+.. autoclass:: numpyro.distributions.continuous.LogNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Normal
+------
+.. autoclass:: numpyro.distributions.continuous.Normal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Pareto
+------
+.. autoclass:: numpyro.distributions.continuous.Pareto
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+StudentT
+--------
+.. autoclass:: numpyro.distributions.continuous.StudentT
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+TruncatedCauchy
+---------------
+.. autoclass:: numpyro.distributions.continuous.TruncatedCauchy
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+TruncatedNormal
+---------------
+.. autoclass:: numpyro.distributions.continuous.TruncatedNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Uniform
+-------
+.. autoclass:: numpyro.distributions.continuous.Uniform
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+
+Discrete Distributions
+======================
+
+Bernoulli
+---------
+.. autofunction:: numpyro.distributions.discrete.Bernoulli
+
+BernoulliLogits
+---------------
+.. autoclass:: numpyro.distributions.discrete.BernoulliLogits
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+BernoulliProbs
+--------------
+.. autoclass:: numpyro.distributions.discrete.BernoulliProbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Binomial
+---------
+.. autofunction:: numpyro.distributions.discrete.Binomial
+
+BinomialLogits
+---------------
+.. autoclass:: numpyro.distributions.discrete.BinomialLogits
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+BinomialProbs
+-------------
+.. autoclass:: numpyro.distributions.discrete.BinomialProbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Categorical
+-----------
+.. autofunction:: numpyro.distributions.discrete.Categorical
+
+CategoricalLogits
+-----------------
+.. autoclass:: numpyro.distributions.discrete.CategoricalLogits
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+CategoricalProbs
+-----------------
+.. autoclass:: numpyro.distributions.discrete.CategoricalProbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Multinomial
+-----------
+.. autofunction:: numpyro.distributions.discrete.Multinomial
+
+MultinomialLogits
+-----------------
+.. autoclass:: numpyro.distributions.discrete.MultinomialLogits
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+MultinomialProbs
+-----------------
+.. autoclass:: numpyro.distributions.discrete.MultinomialProbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Poisson
+-------
+.. autoclass:: numpyro.distributions.discrete.Poisson
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

docs/source/handlers.rst

@@ -1,7 +1,39 @@
-Handlers
-========
+Pyro Primitives
+---------------
+
+.. autofunction:: numpyro.handlers.param
+.. autofunction:: numpyro.handlers.sample
+
+Effect Handlers
+---------------
 
 .. automodule:: numpyro.handlers
+
+.. autoclass:: numpyro.handlers.block
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: numpyro.handlers.replay
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: numpyro.handlers.seed
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: numpyro.handlers.substitute
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: numpyro.handlers.trace
     :members:
     :undoc-members:
     :show-inheritance:

docs/source/index.rst

@@ -16,25 +16,25 @@ Numpyro documentation
 .. toctree::
    :glob:
    :maxdepth: 2
-   :caption: Diagnostics:
+   :caption: Distributions:
 
-   diagnostics
+   distributions
 
 
 .. toctree::
    :glob:
    :maxdepth: 2
-   :caption: Distributions:
+   :caption: Inference Primitives:
 
-   distributions
+   handlers
 
 
 .. toctree::
    :glob:
    :maxdepth: 2
-   :caption: Utilities:
+   :caption: Diagnostics:
 
-   handlers
+   diagnostics
 
 
 Indices and tables

docs/source/mcmc.rst

@@ -1,6 +1,9 @@
 Markov Chain Monte Carlo (MCMC)
 ===============================
 
+Hamiltonian Monte Carlo
+-----------------------
+
 .. autofunction:: numpyro.mcmc.mcmc
 
 .. autofunction:: numpyro.mcmc.hmc
@@ -16,3 +19,8 @@ MCMC Utilities
 --------------
 
 .. autofunction:: numpyro.hmc_util.initialize_model
+
+.. autofunction:: numpyro.util.fori_collect
+
+.. autofunction:: numpyro.diagnostics.summary
+   :noindex:

docs/source/svi.rst

@@ -8,3 +8,7 @@ Stochastic Variational Inference (SVI)
 .. autofunction:: numpyro.svi.svi.update_fn
 
 .. autofunction:: numpyro.svi.svi.evaluate
+
+ELBo
+----
+.. autofunction:: numpyro.svi.elbo