Return scalar variables instead of 1D

Author: ricardoV94

pymc/sampling.py

@@ -2143,17 +2143,19 @@ def draw(
             assert draws[2].shape == (num_draws, 5)
     """
 
-    if not isinstance(vars, (list, tuple)):
-        vars = [vars]
-
     draw_fn = compile_pymc(inputs=[], outputs=vars, mode=mode, **kwargs)
-    drawn_values = zip(*(draw_fn() for _ in range(draws)))
-    drawn_values = [np.stack(v) for v in drawn_values]
 
-    # If only one variable, return the numpy array instead of a list of numpy arrays
     if draws == 1:
-        return drawn_values[0]
-    return drawn_values
+        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(

pymc/tests/test_sampling.py

@@ -1215,44 +1215,56 @@ def test_sample_deterministic():
 
 
 class TestDraw(SeededTest):
-    def test_draw_one_variable(self):
+    def test_univariate(self):
         with pm.Model():
             x = pm.Normal("x")
 
         x_draws = pm.draw(x)
-        assert x_draws.shape == (1,)
+        assert x_draws.shape == ()
 
-    def test_draw_several_variables(self):
+        (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 = 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 samples of a tuple variables
-        draws = pm.draw((x, y, z), draws=num_draws)
+        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)
-
-    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=100)
-        assert mln_draws.shape == (100, 4)
+        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)
-        # Check if the draw function will draw different samples each time
         assert not np.all(np.isclose(x_draws_1, x_draws_2))