add missing docs and remove some TODOs (#529)

Author: fehiepsi

docs/source/autoguide.rst

@@ -11,6 +11,14 @@ AutoContinuous
     :show-inheritance:
     :member-order: bysource
 
+AutoBNAFNormal
+--------------
+.. autoclass:: numpyro.contrib.autoguide.AutoBNAFNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 AutoDiagonalNormal
 ------------------
 .. autoclass:: numpyro.contrib.autoguide.AutoDiagonalNormal

docs/source/distributions.rst

@@ -563,3 +563,11 @@ InverseAutoregressiveTransform
     :undoc-members:
     :show-inheritance:
     :member-order: bysource
+
+BlockNeuralAutoregressiveTransform
+----------------------------------
+.. autoclass:: numpyro.distributions.flows.BlockNeuralAutoregressiveTransform
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

docs/source/handlers.rst

@@ -19,6 +19,14 @@ condition
     :show-inheritance:
     :member-order: bysource
 
+mask
+----
+.. autoclass:: numpyro.handlers.mask
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 replay
 ------
 .. autoclass:: numpyro.handlers.replay

docs/source/primitives.rst

@@ -15,6 +15,10 @@ plate
 -----
 .. autoclass:: numpyro.primitives.plate
 
+deterministic
+-------------
+.. autofunction:: numpyro.primitives.deterministic
+
 factor
 ------
 .. autofunction:: numpyro.primitives.factor

numpyro/distributions/util.py

@@ -46,7 +46,6 @@ def categorical(key, p, shape=()):
     return _categorical(key, p, shape)
 
 
-# TODO: use this sampler in CategoricalLogits
 # TODO: drop this for the next JAX release, see https://github.com/google/jax/pull/1855
 def categorical_logits(key, logits, shape=()):
     shape = shape or logits.shape[:-1]

numpyro/infer/mcmc.py

@@ -810,6 +810,9 @@ class SA(MCMCKernel):
     subset of approximate posterior samples of size num_chains x num_samples
     instead of num_chains x num_samples x adapt_state_size.
 
+    .. note:: We recommend to use this kernel with `progress_bar=False` in :class:`MCMC`
+        to reduce JAX's dispatch overhead.
+
     **References:**
 
     1. *Sample Adaptive MCMC* (https://papers.nips.cc/paper/9107-sample-adaptive-mcmc),

numpyro/primitives.py

@@ -151,7 +151,7 @@ def deterministic(name, value):
     """
     Used to designate deterministic sites in the model. Note that most effect
     handlers will not operate on deterministic sites (except
-    :function:`~numpyro.handlers.trace`), so deterministic sites should be
+    :func:`~numpyro.handlers.trace`), so deterministic sites should be
     side-effect free. The use case for deterministic nodes is to record any
     values in the model execution trace.
 

test/test_hmc_util.py

@@ -394,9 +394,6 @@ def fn(vv_state):
         assert tree.num_proposals > 10
 
 
-# TODO: raise this warning issue upstream, the issue is at this line
-# https://github.com/google/jax/blob/master/jax/numpy/lax_numpy.py#L2732
-@pytest.mark.filterwarnings('ignore:Explicitly requested dtype float64')
 @pytest.mark.parametrize('method', [consensus, parametric_draws])
 @pytest.mark.parametrize('diagonal', [True, False])
 def test_gaussian_subposterior(method, diagonal):