Refactor Multivariate RandomWalk distributions

pymc/distributions/distribution.py

@@ -281,7 +281,7 @@ def __new__(
 
         # Preference is given to size or shape. If not specified, we rely on dims and
         # finally, observed, to determine the shape of the variable.
-        if not ("size" in kwargs or "shape" in kwargs):
+        if kwargs.get("size") is None and kwargs.get("shape") is None:
             if dims is not None:
                 kwargs["shape"] = shape_from_dims(dims, model)
             elif observed is not None:

pymc/distributions/multivariate.py

@@ -381,13 +381,14 @@ class MvStudentT(Continuous):
 
     @classmethod
     def dist(cls, nu, Sigma=None, mu=None, scale=None, tau=None, chol=None, lower=True, **kwargs):
-        if kwargs.get("cov") is not None:
+        cov = kwargs.pop("cov", None)
+        if cov is not None:
             warnings.warn(
                 "Use the scale argument to specify the scale matrix. "
                 "cov will be removed in future versions.",
                 FutureWarning,
             )
-            scale = kwargs.pop("cov")
+            scale = cov
         if Sigma is not None:
             if scale is not None:
                 raise ValueError("Specify only one of scale and Sigma")

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