Create helper pm.draw() to take draws for a given variable (#5340)

Author: danhphan

RELEASE-NOTES.md

@@ -31,10 +31,10 @@ All of the above apply to:
 
 - ⚠ The library is now named, installed and imported as "pymc". For example: `pip install pymc`.
 - ⚠ Theano-PyMC has been replaced with Aesara, so all external references to `theano`, `tt`, and `pymc3.theanof` need to be replaced with `aesara`, `at`, and `pymc.aesaraf` (see [4471](https://github.com/pymc-devs/pymc/pull/4471)).
-- `pm.Distribution(...).logp(x)` is now `pm.logp(pm.Distribution(...), x)`
-- `pm.Distribution(...).logcdf(x)` is now `pm.logcdf(pm.Distribution(...), x)`
-- `pm.Distribution(...).random()` is now `pm.Distribution(...).eval()`
-- `pm.draw_values(...)` and `pm.generate_samples(...)` were removed. The tensors can now be evaluated with `.eval()`.
+- `pm.Distribution(...).logp(x)` is now `pm.logp(pm.Distribution(...), x)`.
+- `pm.Distribution(...).logcdf(x)` is now `pm.logcdf(pm.Distribution(...), x)`.
+- `pm.Distribution(...).random(size=x)` is now `pm.draw(pm.Distribution(...), draws=x)`.
+- `pm.draw_values(...)` and `pm.generate_samples(...)` were removed.
 - `pm.fast_sample_posterior_predictive` was removed.
 - `pm.sample_prior_predictive`, `pm.sample_posterior_predictive` and `pm.sample_posterior_predictive_w` now return an `InferenceData` object by default, instead of a dictionary (see [#5073](https://github.com/pymc-devs/pymc/pull/5073)).
 - `pm.sample_prior_predictive` no longer returns transformed variable values by default. Pass them by name in `var_names` if you want to obtain these draws (see [4769](https://github.com/pymc-devs/pymc/pull/4769)).

pymc/aesaraf.py

@@ -967,7 +967,7 @@ def compile_pymc(inputs, outputs, mode=None, **kwargs):
 
     # Set the default update of a NoDistribution RNG so that it is automatically
     # updated after every function call
-    output_to_list = outputs if isinstance(outputs, list) else [outputs]
+    output_to_list = outputs if isinstance(outputs, (list, tuple)) else [outputs]
     for rv in (
         node
         for node in walk_model(output_to_list, walk_past_rvs=True)

pymc/sampling.py

@@ -42,7 +42,7 @@
 import xarray
 
 from aesara.compile.mode import Mode
-from aesara.graph.basic import Constant
+from aesara.graph.basic import Constant, Variable
 from aesara.tensor.sharedvar import SharedVariable
 from arviz import InferenceData
 from fastprogress.fastprogress import progress_bar
@@ -96,6 +96,7 @@
     "sample_posterior_predictive_w",
     "init_nuts",
     "sample_prior_predictive",
+    "draw",
 ]
 
 STEP_METHODS = (
@@ -2093,6 +2094,70 @@ def sample_prior_predictive(
     return pm.to_inference_data(prior=prior, **ikwargs)
 
 
+def draw(
+    vars: Union[Variable, Sequence[Variable]],
+    draws: int = 1,
+    mode: Optional[Union[str, Mode]] = None,
+    **kwargs,
+) -> Union[np.ndarray, List[np.ndarray]]:
+    """Draw samples for one variable or a list of variables
+
+    Parameters
+    ----------
+    vars
+        A variable or a list of variables for which to draw samples.
+    draws : int
+        Number of samples needed to draw. Detaults to 500.
+    mode
+        The mode used by ``aesara.function`` to compile the graph.
+    **kwargs
+        Keyword arguments for :func:`pymc.aesara.compile_pymc`
+
+    Returns
+    -------
+    List[np.ndarray]
+        A list of numpy arrays.
+
+    Examples
+    --------
+        .. code-block:: python
+
+            import pymc as pm
+
+            # Draw samples for one variable
+            with pm.Model():
+                x = pm.Normal("x")
+            x_draws = pm.draw(x, draws=100)
+            print(x_draws.shape)
+
+            # Draw 1000 samples for several variables
+            with pm.Model():
+                x = pm.Normal("x")
+                y = pm.Normal("y", shape=10)
+                z = pm.Uniform("z", shape=5)
+            num_draws = 1000
+            # Draw samples of a list variables
+            draws = pm.draw([x, y, z], draws=num_draws)
+            assert draws[0].shape == (num_draws,)
+            assert draws[1].shape == (num_draws, 10)
+            assert draws[2].shape == (num_draws, 5)
+    """
+
+    draw_fn = compile_pymc(inputs=[], outputs=vars, mode=mode, **kwargs)
+
+    if draws == 1:
+        return draw_fn()
+
+    # Single variable output
+    if not isinstance(vars, (list, tuple)):
+        drawn_values = (draw_fn() for _ in range(draws))
+        return np.stack(drawn_values)
+
+    # Multiple variable output
+    drawn_values = zip(*(draw_fn() for _ in range(draws)))
+    return [np.stack(v) for v in drawn_values]
+
+
 def _init_jitter(
     model: Model,
     initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]],

pymc/tests/test_sampling.py

@@ -1212,3 +1212,59 @@ def test_sample_deterministic():
         idata = pm.sample(chains=1, draws=50, compute_convergence_checks=False)
 
     np.testing.assert_allclose(idata.posterior["y"], idata.posterior["x"] + 100)
+
+
+class TestDraw(SeededTest):
+    def test_univariate(self):
+        with pm.Model():
+            x = pm.Normal("x")
+
+        x_draws = pm.draw(x)
+        assert x_draws.shape == ()
+
+        (x_draws,) = pm.draw([x])
+        assert x_draws.shape == ()
+
+        x_draws = pm.draw(x, draws=10)
+        assert x_draws.shape == (10,)
+
+        (x_draws,) = pm.draw([x], draws=10)
+        assert x_draws.shape == (10,)
+
+    def test_multivariate(self):
+        with pm.Model():
+            mln = pm.Multinomial("mln", n=5, p=np.array([0.25, 0.25, 0.25, 0.25]))
+
+        mln_draws = pm.draw(mln, draws=1)
+        assert mln_draws.shape == (4,)
+
+        (mln_draws,) = pm.draw([mln], draws=1)
+        assert mln_draws.shape == (4,)
+
+        mln_draws = pm.draw(mln, draws=10)
+        assert mln_draws.shape == (10, 4)
+
+        (mln_draws,) = pm.draw([mln], draws=10)
+        assert mln_draws.shape == (10, 4)
+
+    def test_multiple_variables(self):
+        with pm.Model():
+            x = pm.Normal("x")
+            y = pm.Normal("y", shape=10)
+            z = pm.Uniform("z", shape=5)
+            w = pm.Dirichlet("w", a=[1, 1, 1])
+
+        num_draws = 100
+        draws = pm.draw((x, y, z, w), draws=num_draws)
+        assert draws[0].shape == (num_draws,)
+        assert draws[1].shape == (num_draws, 10)
+        assert draws[2].shape == (num_draws, 5)
+        assert draws[3].shape == (num_draws, 3)
+
+    def test_draw_different_samples(self):
+        with pm.Model():
+            x = pm.Normal("x")
+
+        x_draws_1 = pm.draw(x, 100)
+        x_draws_2 = pm.draw(x, 100)
+        assert not np.all(np.isclose(x_draws_1, x_draws_2))