Fix get_support_shape

Author: AlexAndorra

pymc/distributions/multivariate.py

@@ -2487,7 +2487,7 @@ def dist(cls, sigma=1, zerosum_axes=None, support_shape=None, **kwargs):
 
         if support_shape is None:
             if zerosum_axes > 0:
-                raise ValueError("You must specify shape or support_shape parameter")
+                raise ValueError("You must specify dims, shape or support_shape parameter")
             # edge-case doesn't work for now, because at.stack in get_support_shape fails
             # else:
             #     support_shape = () # because it's just a Normal in that case
@@ -2553,7 +2553,6 @@ def rv_op(cls, sigma, zerosum_axes, support_shape, size=None):
         )(normal_dist, sigma, support_shape)
 
         # TODO:
-        # refactor ZSN tests
         # test get_support_shape with 2D
         # test ZSN logp
         # test ZSN variance

pymc/distributions/shape_utils.py

@@ -741,8 +741,13 @@ def get_support_shape(
         inferred_support_shape = support_shape
     # If there are two sources of information for the support shapes, assert they are consistent:
     elif support_shape is not None:
-        inferred_support_shape = Assert(msg="support_shape does not match last shape dimension")(
-            inferred_support_shape, at.all(at.eq(inferred_support_shape, support_shape))
+        inferred_support_shape = at.stack(
+            [
+                Assert(msg="support_shape does not match last shape dimension")(
+                    inferred, at.eq(inferred, explicit)
+                )
+                for inferred, explicit in zip(inferred_support_shape, support_shape)
+            ]
         )
 
     return inferred_support_shape