Fix scalar size issue in CategoricalRV

aesara-devs · Mar 26, 2021 · 0237886 · 0237886
1 parent ce9c17f
commit 0237886
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/aesara/tensor/random/basic.py b/aesara/tensor/random/basic.py
@@ -413,10 +413,14 @@ def rng_fn(cls, rng, p, size):
         size = tuple(np.atleast_1d(size))
         ind_shape = p.shape[:-1]
 
-        if len(size) > 0 and size[-len(ind_shape) :] != ind_shape:
-            raise ValueError("Parameters shape and size do not match.")
+        if len(ind_shape) > 0:
+            if len(size) > 0 and size[-len(ind_shape) :] != ind_shape:
+                raise ValueError("Parameters shape and size do not match.")
+
+            samples_shape = size[: -len(ind_shape)] + ind_shape
+        else:
+            samples_shape = size
 
-        samples_shape = size[: -len(ind_shape)] + ind_shape
         unif_samples = rng.uniform(size=samples_shape)
         samples = vsearchsorted(p.cumsum(axis=-1), unif_samples)
 

diff --git a/tests/tensor/random/test_basic.py b/tests/tensor/random/test_basic.py
@@ -589,6 +589,10 @@ def test_categorical_samples():
 
     rng_state = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(1234)))
 
+    assert categorical.rng_fn(rng_state, np.array([1.0 / 3.0] * 3), size=10).shape == (
+        10,
+    )
+
     p = np.array([[100000, 1, 1], [1, 100000, 1], [1, 1, 100000]], dtype=config.floatX)
     p = p / p.sum(axis=-1)