Fix AutoNormal.quantiles (#1066)

* fix AutoNormal.quantiles and add docs

* force tfp release less than 0.14.0.dev

* rearange docs

* fix docs

Author: fehiepsi

docs/source/autoguide.rst

@@ -3,6 +3,14 @@ Automatic Guide Generation
 
 .. automodule:: numpyro.infer.autoguide
 
+AutoGuide
+---------
+.. autoclass:: numpyro.infer.autoguide.AutoGuide
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 AutoContinuous
 --------------
 .. autoclass:: numpyro.infer.autoguide.AutoContinuous
@@ -73,4 +81,4 @@ AutoDelta
     :members:
     :undoc-members:
     :show-inheritance:
-    :member-order: bysource
+    :member-order: bysource

numpyro/infer/autoguide.py

@@ -113,15 +113,18 @@ def __call__(self, *args, **kwargs):
         raise NotImplementedError
 
     @abstractmethod
-    def sample_posterior(self, rng_key, params, *args, **kwargs):
+    def sample_posterior(self, rng_key, params, sample_shape=()):
         """
         Generate samples from the approximate posterior over the latent
         sites in the model.
 
-        :param jax.random.PRNGKey rng_key: PRNG seed.
-        :param params: Current parameters of model and autoguide.
-        :param sample_shape: (keyword argument) shape of samples to be drawn.
-        :return: batch of samples from the approximate posterior.
+        :param jax.random.PRNGKey rng_key: random key to be used draw samples.
+        :param dict params: Current parameters of model and autoguide.
+            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
+            method from :class:`~numpyro.infer.svi.SVI`.
+        :param tuple sample_shape: sample shape of each latent site, defaults to ().
+        :return: a dict containing samples drawn the this guide.
+        :rtype: dict
         """
         raise NotImplementedError
 
@@ -161,14 +164,41 @@ def _setup_prototype(self, *args, **kwargs):
             elif site["type"] == "plate":
                 self._prototype_frame_full_sizes[name] = site["args"][0]
 
+    def median(self, params):
+        """
+        Returns the posterior median value of each latent variable.
+
+        :param dict params: A dict containing parameter values.
+            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
+            method from :class:`~numpyro.infer.svi.SVI`.
+        :return: A dict mapping sample site name to median value.
+        :rtype: dict
+        """
+        raise NotImplementedError
+
+    def quantiles(self, params, quantiles):
+        """
+        Returns posterior quantiles each latent variable. Example::
+
+            print(guide.quantiles(params, [0.05, 0.5, 0.95]))
+
+        :param dict params: A dict containing parameter values.
+            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
+            method from :class:`~numpyro.infer.svi.SVI`.
+        :param list quantiles: A list of requested quantiles between 0 and 1.
+        :return: A dict mapping sample site name to an array of quantile values.
+        :rtype: dict
+        """
+        raise NotImplementedError
+
 
 class AutoNormal(AutoGuide):
     """
     This implementation of :class:`AutoGuide` uses Normal distributions
     to construct a guide over the entire latent space. The guide does not
     depend on the model's ``*args, **kwargs``.
 
-    This should be equivalent to :class: `AutoDiagonalNormal` , but with
+    This should be equivalent to :class:`AutoDiagonalNormal` , but with
     more convenient site names and with better support for mean field ELBO.
 
     Usage::
@@ -231,12 +261,6 @@ def _setup_prototype(self, *args, **kwargs):
                     )
 
     def __call__(self, *args, **kwargs):
-        """
-        An automatic guide with the same ``*args, **kwargs`` as the base ``model``.
-
-        :return: A dict mapping sample site name to sampled value.
-        :rtype: dict
-        """
         if self.prototype_trace is None:
             # run model to inspect the model structure
             self._setup_prototype(*args, **kwargs)
@@ -314,10 +338,15 @@ def median(self, params):
         return self._constrain(locs)
 
     def quantiles(self, params, quantiles):
-        quantiles = jnp.array(quantiles)[..., None]
+        quantiles = jnp.array(quantiles)
         locs = {k: params["{}_{}_loc".format(k, self.prefix)] for k in self._init_locs}
         scales = {k: params["{}_{}_scale".format(k, self.prefix)] for k in locs}
-        latent = {k: dist.Normal(locs[k], scales[k]).icdf(quantiles) for k in locs}
+        latent = {
+            k: dist.Normal(locs[k], scales[k]).icdf(
+                quantiles.reshape((-1,) + (1,) * jnp.ndim(locs[k]))
+            )
+            for k in locs
+        }
         return self._constrain(latent)
 
 
@@ -413,12 +442,6 @@ def sample_posterior(self, rng_key, params, sample_shape=()):
         return latent_samples
 
     def median(self, params):
-        """
-        Returns the posterior median value of each latent variable.
-
-        :return: A dict mapping sample site name to median tensor.
-        :rtype: dict
-        """
         locs = {k: params["{}_{}_loc".format(k, self.prefix)] for k in self._init_locs}
         return locs
 
@@ -473,12 +496,6 @@ def _sample_latent(self, *args, **kwargs):
         )
 
     def __call__(self, *args, **kwargs):
-        """
-        An automatic guide with the same ``*args, **kwargs`` as the base ``model``.
-
-        :return: A dict mapping sample site name to sampled value.
-        :rtype: dict
-        """
         if self.prototype_trace is None:
             # run model to inspect the model structure
             self._setup_prototype(*args, **kwargs)
@@ -585,49 +602,11 @@ def get_posterior(self, params):
         return dist.TransformedDistribution(base_dist, transform)
 
     def sample_posterior(self, rng_key, params, sample_shape=()):
-        """
-        Get samples from the learned posterior.
-
-        :param jax.random.PRNGKey rng_key: random key to be used draw samples.
-        :param dict params: Current parameters of model and autoguide.
-            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
-            method from :class:`~numpyro.infer.svi.SVI`.
-        :param tuple sample_shape: batch shape of each latent sample, defaults to ().
-        :return: a dict containing samples drawn the this guide.
-        :rtype: dict
-        """
         latent_sample = handlers.substitute(
             handlers.seed(self._sample_latent, rng_key), params
         )(sample_shape=sample_shape)
         return self._unpack_and_constrain(latent_sample, params)
 
-    def median(self, params):
-        """
-        Returns the posterior median value of each latent variable.
-
-        :param dict params: A dict containing parameter values.
-            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
-            method from :class:`~numpyro.infer.svi.SVI`.
-        :return: A dict mapping sample site name to median tensor.
-        :rtype: dict
-        """
-        raise NotImplementedError
-
-    def quantiles(self, params, quantiles):
-        """
-        Returns posterior quantiles each latent variable. Example::
-
-            print(guide.quantiles(opt_state, [0.05, 0.5, 0.95]))
-
-        :param dict params: A dict containing parameter values.
-            The parameters can be obtained using :meth:`~numpyro.infer.svi.SVI.get_params`
-            method from :class:`~numpyro.infer.svi.SVI`.
-        :param list quantiles: A list of requested quantiles between 0 and 1.
-        :return: A dict mapping sample site name to a list of quantile values.
-        :rtype: dict
-        """
-        raise NotImplementedError
-
 
 class AutoDiagonalNormal(AutoContinuous):
     """

test/infer/test_autoguide.py

@@ -81,6 +81,10 @@ def body_fn(i, val):
     )
     assert_allclose(jnp.mean(posterior_samples["beta"], 0), true_coefs, atol=0.05)
 
+    if auto_class not in [AutoDelta, AutoIAFNormal, AutoBNAFNormal]:
+        quantiles = guide.quantiles(params, [0.2, 0.5, 0.8])
+        assert quantiles["beta"].shape == (3, 2)
+
     # Predictive can be instantiated from posterior samples...
     predictive = Predictive(model, posterior_samples=posterior_samples)
     predictive_samples = predictive(random.PRNGKey(1), None)