Refactor get_support_shape

* Do not call `at.stack` until the end
* Allow to pass single shape_offset

Author: ricardoV94

pymc/distributions/shape_utils.py

@@ -20,7 +20,7 @@
 import warnings
 
 from functools import singledispatch
-from typing import Any, Optional, Sequence, Tuple, Union
+from typing import Any, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
@@ -671,7 +671,7 @@ def get_support_shape(
     observed: Optional[Any] = None,
     support_shape_offset: Sequence[int] = None,
     ndim_supp: int = 1,
-):
+) -> Optional[TensorVariable]:
     """Extract the support shapes from shape / dims / observed information
 
     Parameters
@@ -702,46 +702,61 @@ def get_support_shape(
         raise NotImplementedError("ndim_supp must be bigger than 0")
     if support_shape_offset is None:
         support_shape_offset = [0] * ndim_supp
-    inferred_support_shape = None
+    elif isinstance(support_shape_offset, int):
+        support_shape_offset = [support_shape_offset] * ndim_supp
+    inferred_support_shape: Optional[Sequence[Union[int, np.ndarray, Variable]]] = None
 
     if shape is not None:
         shape = to_tuple(shape)
         assert isinstance(shape, tuple)
-        inferred_support_shape = at.stack(
-            [shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)]
-        )
+        if len(shape) < ndim_supp:
+            raise ValueError(
+                f"Number of shape dimensions is too small for ndim_supp of {ndim_supp}"
+            )
+        inferred_support_shape = [
+            shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)
+        ]
 
     if inferred_support_shape is None and dims is not None:
         dims = convert_dims(dims)
         assert isinstance(dims, tuple)
+        if len(dims) < ndim_supp:
+            raise ValueError(f"Number of dims is too small for ndim_supp of {ndim_supp}")
         model = modelcontext(None)
-        inferred_support_shape = at.stack(
-            [
-                model.dim_lengths[dims[i]] - support_shape_offset[i]  # type: ignore
-                for i in np.arange(-ndim_supp, 0)
-            ]
-        )
+        inferred_support_shape = [
+            model.dim_lengths[dims[i]] - support_shape_offset[i]  # type: ignore
+            for i in np.arange(-ndim_supp, 0)
+        ]
 
     if inferred_support_shape is None and observed is not None:
         observed = convert_observed_data(observed)
-        inferred_support_shape = at.stack(
-            [observed.shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)]
-        )
+        if observed.ndim < ndim_supp:
+            raise ValueError(
+                f"Number of observed dimensions is too small for ndim_supp of {ndim_supp}"
+            )
+        inferred_support_shape = [
+            observed.shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)
+        ]
 
-    if inferred_support_shape is None:
+    # We did not learn anything
+    if inferred_support_shape is None and support_shape is None:
+        return None
+    # Only source of information was the originally provided support_shape
+    elif inferred_support_shape is None:
         inferred_support_shape = support_shape
-    # If there are two sources of information for the support shapes, assert they are consistent:
+    # There were two sources of support_shape, make sure they are consistent
     elif support_shape is not None:
-        inferred_support_shape = at.stack(
-            [
+        inferred_support_shape = [
+            cast(
+                Variable,
                 Assert(msg="support_shape does not match respective shape dimension")(
                     inferred, at.eq(inferred, explicit)
-                )
-                for inferred, explicit in zip(inferred_support_shape, support_shape)
-            ]
-        )
+                ),
+            )
+            for inferred, explicit in zip(inferred_support_shape, support_shape)
+        ]
 
-    return inferred_support_shape
+    return at.stack(inferred_support_shape)
 
 
 def get_support_shape_1d(
@@ -751,21 +766,18 @@ def get_support_shape_1d(
     dims: Optional[Dims] = None,
     observed: Optional[Any] = None,
     support_shape_offset: int = 0,
-):
+) -> Optional[TensorVariable]:
     """Helper function for cases when you just care about one dimension."""
-    if support_shape is not None:
-        support_shape_tuple = (support_shape,)
-    else:
-        support_shape_tuple = None
-
     support_shape_tuple = get_support_shape(
-        support_shape_tuple,
+        support_shape=(support_shape,) if support_shape is not None else None,
         shape=shape,
         dims=dims,
         observed=observed,
         support_shape_offset=(support_shape_offset,),
     )
-    if support_shape_tuple is not None:
-        (support_shape,) = support_shape_tuple
 
-    return support_shape
+    if support_shape_tuple is not None:
+        (support_shape_,) = support_shape_tuple
+        return support_shape_
+    else:
+        return None

pymc/tests/distributions/test_multivariate.py

@@ -1466,7 +1466,13 @@ def test_zsn_shape(self, zerosum_axes):
     @pytest.mark.parametrize(
         "error, match, shape, support_shape, zerosum_axes",
         [
-            (IndexError, "index out of range", (3, 4, 5), None, 4),
+            (
+                ValueError,
+                "Number of shape dimensions is too small for ndim_supp of 4",
+                (3, 4, 5),
+                None,
+                4,
+            ),
             (AssertionError, "does not match", (3, 4), (3,), None),  # support_shape should be 4
             (
                 AssertionError,