Refactor get_steps to work with multivariate support shapes

Author: AlexAndorra

pymc/distributions/shape_utils.py

@@ -20,19 +20,23 @@
 import warnings
 
 from functools import singledispatch
-from typing import Optional, Sequence, Tuple, Union
+from typing import Any, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
 from aesara import config
 from aesara import tensor as at
 from aesara.graph.basic import Variable
 from aesara.graph.op import Op, compute_test_value
+from aesara.raise_op import Assert
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
+from pymc.aesaraf import convert_observed_data
+from pymc.model import modelcontext
+
 __all__ = [
     "to_tuple",
     "shapes_broadcasting",
@@ -666,3 +670,100 @@ def change_specify_shape_size(op, ss, new_size, expand) -> TensorVariable:
 
     # specify_shape has a wrong signature https://github.com/aesara-devs/aesara/issues/1164
     return at.specify_shape(new_var, new_shapes)  # type: ignore
+
+
+def get_support_shape(
+    support_shape: Optional[Sequence[Union[int, np.ndarray, TensorVariable]]],
+    *,
+    shape: Optional[Shape] = None,
+    dims: Optional[Dims] = None,
+    observed: Optional[Any] = None,
+    support_shape_offset: Sequence[int] = None,
+    ndim_supp: int = 1,
+):
+    """Extract length of support shapes from shape / dims / observed information
+
+    Parameters
+    ----------
+    support_shape:
+        User-specified support shape for multivariate distribution
+    shape:
+        User-specified shape for multivariate distribution
+    dims:
+        User-specified dims for multivariate distribution
+    observed:
+        User-specified observed data from multivariate distribution
+    support_shape_offset:
+        Difference between last shape dimensions and the length of explicit support shapes in multivariate distribution, defaults to 0.
+        For timeseries, this is shape[-1] = support_shape[-1] + 1
+    ndim_supp:
+        Number of support dimensions of the given multivariate distribution, defaults to 1
+
+    Returns
+    -------
+    support_shape
+        Support shape, if specified directly by user, or inferred from the last dimensions of
+        shape / dims / observed. When two sources of support shape information are provided,
+        a symbolic Assert is added to ensure they are consistent.
+    """
+    if support_shape_offset is None:
+        support_shape_offset = [0] * ndim_supp
+    inferred_support_shape = None
+
+    if shape is not None:
+        shape = to_tuple(shape)
+        assert isinstance(shape, tuple)
+        inferred_support_shape = at.stack(
+            [shape[-i - 1] - support_shape_offset[-i - 1] for i in range(ndim_supp)]
+        )
+
+    if inferred_support_shape is None and dims is not None:
+        dims = convert_dims(dims)
+        assert isinstance(dims, tuple)
+        model = modelcontext(None)
+        inferred_support_shape = at.stack(
+            [
+                model.dim_lengths[dims[-i - 1]] - support_shape_offset[-i - 1]  # type: ignore
+                for i in range(ndim_supp)
+            ]
+        )
+
+    if inferred_support_shape is None and observed is not None:
+        observed = convert_observed_data(observed)
+        inferred_support_shape = at.stack(
+            [observed.shape[-i - 1] - support_shape_offset[-i - 1] for i in range(ndim_supp)]
+        )
+
+    if 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:
+    elif support_shape is not None:
+        inferred_support_shape = Assert(msg="Steps do not match last shape dimension")(
+            inferred_support_shape, at.all(at.eq(inferred_support_shape, support_shape))
+        )
+    return inferred_support_shape
+
+
+def get_support_shape_1d(
+    support_shape: Optional[Union[int, np.ndarray, TensorVariable]],
+    *,
+    shape: Optional[Shape] = None,
+    dims: Optional[Dims] = None,
+    observed: Optional[Any] = None,
+    support_shape_offset: int = 0,
+):
+    """Helper function for cases when you just care about one dimension."""
+    if support_shape is not None:
+        support_shape_tuple = (support_shape,)
+
+    support_shape_tuple = get_support_shape(
+        support_shape_tuple,
+        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

pymc/distributions/timeseries.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 import warnings
 
-from typing import Any, Optional, Union
+from typing import Optional
 
 import aesara
 import aesara.tensor as at
@@ -24,12 +24,11 @@
 from aesara import scan
 from aesara.graph import FunctionGraph, rewrite_graph
 from aesara.graph.basic import Node, clone_replace
-from aesara.raise_op import Assert
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.rewriting.basic import ShapeFeature, topo_constant_folding
 
-from pymc.aesaraf import convert_observed_data, floatX, intX
+from pymc.aesaraf import floatX, intX
 from pymc.distributions import distribution, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.distribution import (
@@ -40,14 +39,11 @@
 )
 from pymc.distributions.logprob import ignore_logprob, logp
 from pymc.distributions.shape_utils import (
-    Dims,
-    Shape,
     _change_dist_size,
     change_dist_size,
-    convert_dims,
+    get_support_shape_1d,
     to_tuple,
 )
-from pymc.model import modelcontext
 from pymc.util import check_dist_not_registered
 
 __all__ = [
@@ -61,61 +57,6 @@
 ]
 
 
-def get_steps(
-    steps: Optional[Union[int, np.ndarray, TensorVariable]],
-    *,
-    shape: Optional[Shape] = None,
-    dims: Optional[Dims] = None,
-    observed: Optional[Any] = None,
-    step_shape_offset: int = 0,
-):
-    """Extract number of steps from shape / dims / observed information
-
-    Parameters
-    ----------
-    steps:
-        User specified steps for timeseries distribution
-    shape:
-        User specified shape for timeseries distribution
-    dims:
-        User specified dims for timeseries distribution
-    observed:
-        User specified observed data from timeseries distribution
-    step_shape_offset:
-        Difference between last shape dimension and number of steps in timeseries
-        distribution, defaults to 0
-
-    Returns
-    -------
-    steps
-        Steps, if specified directly by user, or inferred from the last dimension of
-        shape / dims / observed. When two sources of step information are provided,
-        a symbolic Assert is added to ensure they are consistent.
-    """
-    inferred_steps = None
-    if shape is not None:
-        shape = to_tuple(shape)
-        inferred_steps = shape[-1] - step_shape_offset
-
-    if inferred_steps is None and dims is not None:
-        dims = convert_dims(dims)
-        model = modelcontext(None)
-        inferred_steps = model.dim_lengths[dims[-1]] - step_shape_offset
-
-    if inferred_steps is None and observed is not None:
-        observed = convert_observed_data(observed)
-        inferred_steps = observed.shape[-1] - step_shape_offset
-
-    if inferred_steps is None:
-        inferred_steps = steps
-    # If there are two sources of information for the steps, assert they are consistent
-    elif steps is not None:
-        inferred_steps = Assert(msg="Steps do not match last shape dimension")(
-            inferred_steps, at.eq(inferred_steps, steps)
-        )
-    return inferred_steps
-
-
 class RandomWalkRV(SymbolicRandomVariable):
     """RandomWalk Variable"""
 
@@ -132,21 +73,21 @@ class RandomWalk(Distribution):
     rv_type = RandomWalkRV
 
     def __new__(cls, *args, steps=None, **kwargs):
-        steps = get_steps(
-            steps=steps,
+        steps = get_support_shape_1d(
+            support_shape=steps,
             shape=None,  # Shape will be checked in `cls.dist`
             dims=kwargs.get("dims", None),
             observed=kwargs.get("observed", None),
-            step_shape_offset=1,
+            support_shape_offset=1,
         )
         return super().__new__(cls, *args, steps=steps, **kwargs)
 
     @classmethod
     def dist(cls, init_dist, innovation_dist, steps=None, **kwargs) -> at.TensorVariable:
-        steps = get_steps(
-            steps=steps,
+        steps = get_support_shape_1d(
+            support_shape=steps,
             shape=kwargs.get("shape"),
-            step_shape_offset=1,
+            support_shape_offset=1,
         )
         if steps is None:
             raise ValueError("Must specify steps or shape parameter")
@@ -391,12 +332,12 @@ class AR(Distribution):
     def __new__(cls, name, rho, *args, steps=None, constant=False, ar_order=None, **kwargs):
         rhos = at.atleast_1d(at.as_tensor_variable(floatX(rho)))
         ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
-        steps = get_steps(
-            steps=steps,
+        steps = get_support_shape_1d(
+            support_shape=steps,
             shape=None,  # Shape will be checked in `cls.dist`
             dims=kwargs.get("dims", None),
             observed=kwargs.get("observed", None),
-            step_shape_offset=ar_order,
+            support_shape_offset=ar_order,
         )
         return super().__new__(
             cls, name, rhos, *args, steps=steps, constant=constant, ar_order=ar_order, **kwargs
@@ -427,7 +368,9 @@ def dist(
             init_dist = kwargs.pop("init")
 
         ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order)
-        steps = get_steps(steps=steps, shape=kwargs.get("shape", None), step_shape_offset=ar_order)
+        steps = get_support_shape_1d(
+            support_shape=steps, shape=kwargs.get("shape", None), support_shape_offset=ar_order
+        )
         if steps is None:
             raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps), ndim=0)

pymc/tests/distributions/test_shape_utils.py

@@ -18,9 +18,10 @@
 import numpy as np
 import pytest
 
-from aesara import Mode
 from aesara import tensor as at
+from aesara.compile.mode import Mode
 from aesara.graph import Constant, ancestors
+from aesara.tensor import TensorVariable
 from aesara.tensor.random import normal
 from aesara.tensor.shape import SpecifyShape
 
@@ -36,10 +37,12 @@
     convert_shape,
     convert_size,
     get_broadcastable_dist_samples,
+    get_support_shape_1d,
     rv_size_is_none,
     shapes_broadcasting,
     to_tuple,
 )
+from pymc.model import Model
 
 test_shapes = [
     (tuple(), (1,), (4,), (5, 4)),
@@ -622,3 +625,63 @@ def test_change_specify_shape_size_multivariate():
         new_x.eval({batch: 5, supp: 3}).shape == (10, 5, 5, 3)
         with pytest.raises(AssertionError, match=re.escape("expected (None, None, 5, 3)")):
             new_x.eval({batch: 6, supp: 3}).shape == (10, 5, 5, 3)
+
+
+@pytest.mark.parametrize(
+    "steps, shape, step_shape_offset, expected_steps, consistent",
+    [
+        (10, None, 0, 10, True),
+        (10, None, 1, 10, True),
+        (None, (10,), 0, 10, True),
+        (None, (10,), 1, 9, True),
+        (None, (10, 5), 0, 5, True),
+        (None, None, 0, None, True),
+        (10, (10,), 0, 10, True),
+        (10, (11,), 1, 10, True),
+        (10, (5, 5), 0, 5, False),
+        (10, (5, 10), 1, 9, False),
+    ],
+)
+@pytest.mark.parametrize("info_source", ("shape", "dims", "observed"))
+def test_get_support_shape_1d(
+    info_source, steps, shape, step_shape_offset, expected_steps, consistent
+):
+    if info_source == "shape":
+        inferred_steps = get_support_shape_1d(
+            support_shape=steps, shape=shape, support_shape_offset=step_shape_offset
+        )
+
+    elif info_source == "dims":
+        if shape is None:
+            dims = None
+            coords = {}
+        else:
+            dims = tuple(str(i) for i, shape in enumerate(shape))
+            coords = {str(i): range(shape) for i, shape in enumerate(shape)}
+        with Model(coords=coords):
+            inferred_steps = get_support_shape_1d(
+                support_shape=steps, dims=dims, support_shape_offset=step_shape_offset
+            )
+
+    elif info_source == "observed":
+        if shape is None:
+            observed = None
+        else:
+            observed = np.zeros(shape)
+        inferred_steps = get_support_shape_1d(
+            support_shape=steps, observed=observed, support_shape_offset=step_shape_offset
+        )
+
+    if not isinstance(inferred_steps, TensorVariable):
+        assert inferred_steps == expected_steps
+    else:
+        if consistent:
+            assert inferred_steps.eval() == expected_steps
+        else:
+            # check that inferred steps is still correct by ignoring the assert
+            f = aesara.function(
+                [], inferred_steps, mode=Mode().including("local_remove_all_assert")
+            )
+            assert f() == expected_steps
+            with pytest.raises(AssertionError, match="Steps do not match"):
+                inferred_steps.eval()