shape problem when sampling

Author: ferrine

pymc3/tests/test_distributions_random.py

@@ -96,20 +96,22 @@ def __init__(self, *args, **kwargs):
             super(BaseTestCases.BaseTestCase, self).__init__(*args, **kwargs)
             self.model = pm.Model()
 
-        def get_random_variable(self, shape, with_vector_params=False):
+        def get_random_variable(self, shape, with_vector_params=False, name=None):
             if with_vector_params:
                 params = {key: value * np.ones(self.shape, dtype=np.dtype(type(value))) for
                           key, value in self.params.items()}
             else:
                 params = self.params
-            name = self.distribution.__name__
+            if name is None:
+                name = self.distribution.__name__
             with self.model:
                 if shape is None:
                     return self.distribution(name, transform=None, **params)
                 else:
                     return self.distribution(name, shape=shape, transform=None, **params)
 
-        def sample_random_variable(self, random_variable, size):
+        @staticmethod
+        def sample_random_variable(random_variable, size):
             try:
                 return random_variable.random(size=size)
             except AttributeError:
@@ -145,7 +147,7 @@ def test_parameters_1d_shape(self):
                 else:
                     expected = np.atleast_1d(size).tolist()
                 expected.append(self.shape)
-                actual = np.atleast_1d(self.sample_random_variable(rv, size)).shape
+                actual = self.sample_random_variable(rv, size).shape
                 self.assertSequenceEqual(expected, actual)
 
         def test_broadcast_shape(self):
@@ -160,6 +162,26 @@ def test_broadcast_shape(self):
                 actual = np.atleast_1d(self.sample_random_variable(rv, size)).shape
                 self.assertSequenceEqual(expected, actual)
 
+        def test_different_shapes_and_sample_sizes(self):
+            shapes = [(), (1,), (1, 1), (1, 2), (10, 10, 1), (10, 10, 2)]
+            prefix = self.distribution.__name__
+            expected = []
+            actual = []
+            for shape in shapes:
+                rv = self.get_random_variable(shape, name='%s_%s' % (prefix, shape))
+                for size in (None, 1, 5, (4, 5)):
+                    if size is None:
+                        s = []
+                    else:
+                        try:
+                            s = list(size)
+                        except TypeError:
+                            s = [size]
+                    s.extend(shape)
+                    expected.append(tuple(s))
+                    actual.append(self.sample_random_variable(rv, size).shape)
+            self.assertSequenceEqual(expected, actual)
+
 
 class TestNormal(BaseTestCases.BaseTestCase):
     distribution = pm.Normal