Obtain timeseries step information from dims or observed

pymc/aesaraf.py

@@ -81,16 +81,13 @@
     "set_at_rng",
     "at_rng",
     "take_along_axis",
-    "pandas_to_array",
+    "convert_observed_data",
 ]
 
 
-def pandas_to_array(data):
-    """Convert a pandas object to a NumPy array.
+def convert_observed_data(data):
+    """Convert user provided dataset to accepted formats."""
 
-    XXX: When `data` is a generator, this will return an Aesara tensor!
-
-    """
     if hasattr(data, "to_numpy") and hasattr(data, "isnull"):
         # typically, but not limited to pandas objects
         vals = data.to_numpy()

pymc/data.py

@@ -33,7 +33,7 @@
 
 import pymc as pm
 
-from pymc.aesaraf import pandas_to_array
+from pymc.aesaraf import convert_observed_data
 
 __all__ = [
     "get_data",
@@ -636,9 +636,9 @@ def Data(
         )
     name = model.name_for(name)
 
-    # `pandas_to_array` takes care of parameter `value` and
+    # `convert_observed_data` takes care of parameter `value` and
     # transforms it to something digestible for Aesara.
-    arr = pandas_to_array(value)
+    arr = convert_observed_data(value)
 
     if mutable is None:
         major, minor = (int(v) for v in pm.__version__.split(".")[:2])

pymc/distributions/shape_utils.py

@@ -26,7 +26,7 @@
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
-from pymc.aesaraf import pandas_to_array
+from pymc.aesaraf import convert_observed_data
 
 __all__ = [
     "to_tuple",
@@ -558,7 +558,7 @@ def resize_from_observed(
         Observations as numpy array or `Variable`.
     """
     if not hasattr(observed, "shape"):
-        observed = pandas_to_array(observed)
+        observed = convert_observed_data(observed)
     ndim_resize = observed.ndim - ndim_implied
     resize_shape = tuple(observed.shape[d] for d in range(ndim_resize))
     return resize_shape, observed

pymc/distributions/timeseries.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 import warnings
 
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import aesara
 import aesara.tensor as at
@@ -31,13 +31,20 @@
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.utils import normalize_size_param
 
-from pymc.aesaraf import change_rv_size, floatX, intX
+from pymc.aesaraf import change_rv_size, convert_observed_data, floatX, intX
 from pymc.distributions import distribution, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import SymbolicDistribution, _moment, moment
 from pymc.distributions.logprob import ignore_logprob, logp
-from pymc.distributions.shape_utils import Shape, rv_size_is_none, to_tuple
+from pymc.distributions.shape_utils import (
+    Dims,
+    Shape,
+    convert_dims,
+    rv_size_is_none,
+    to_tuple,
+)
+from pymc.model import modelcontext
 from pymc.util import check_dist_not_registered
 
 __all__ = [
@@ -50,51 +57,61 @@
 ]
 
 
-def get_steps_from_shape(
+def get_steps(
     steps: Optional[Union[int, np.ndarray, TensorVariable]],
-    shape: Optional[Shape],
+    *,
+    shape: Optional[Shape] = None,
+    dims: Optional[Dims] = None,
+    observed: Optional[Any] = None,
     step_shape_offset: int = 0,
 ):
-    """Extract number of steps from shape information
+    """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
 
-    Raises
-    ------
-    ValueError
-        If neither shape nor steps are provided
-
     Returns
     -------
     steps
         Steps, if specified directly by user, or inferred from the last dimension of
-        shape. When both steps and shape are provided, a symbolic Assert is added
-        to make sure they are consistent.
+        shape / dims / observed. When two sources of step information are provided,
+        a symbolic Assert is added to ensure they are consistent.
     """
-    steps_from_shape = None
+    inferred_steps = None
     if shape is not None:
         shape = to_tuple(shape)
         if shape[-1] is not ...:
-            steps_from_shape = shape[-1] - step_shape_offset
-    if steps is None:
-        if steps_from_shape is not None:
-            steps = steps_from_shape
-        else:
-            raise ValueError("Must specify steps or shape parameter")
-    elif steps_from_shape is not None:
-        # Assert that steps and shape are consistent
-        steps = Assert(msg="Steps do not match last shape dimension")(
-            steps, at.eq(steps, steps_from_shape)
+            inferred_steps = shape[-1] - step_shape_offset
+
+    if inferred_steps is None and dims is not None:
+        dims = convert_dims(dims)
+        if dims[-1] is not ...:
+            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 steps
+    return inferred_steps
 
 
 class GaussianRandomWalkRV(RandomVariable):
@@ -212,26 +229,38 @@ class GaussianRandomWalk(distribution.Continuous):
 
         .. warning:: init will be cloned, rendering them independent of the ones passed as input.
 
-    steps : int
-        Number of steps in Gaussian Random Walks (steps > 0).
+    steps : int, optional
+        Number of steps in Gaussian Random Walk (steps > 0). Only needed if size is
+        used to specify distribution
     """
 
     rv_op = gaussianrandomwalk
 
-    def __new__(cls, name, mu=0.0, sigma=1.0, init=None, steps=None, **kwargs):
-        if init is not None:
-            check_dist_not_registered(init)
-        return super().__new__(cls, name, mu, sigma, init, steps, **kwargs)
+    def __new__(cls, *args, steps=None, **kwargs):
+        steps = get_steps(
+            steps=steps,
+            shape=None,  # Shape will be checked in `cls.dist`
+            dims=kwargs.get("dims", None),
+            observed=kwargs.get("observed", None),
+            step_shape_offset=1,
+        )
+        return super().__new__(cls, *args, steps=steps, **kwargs)
 
     @classmethod
     def dist(
-        cls, mu=0.0, sigma=1.0, init=None, steps=None, size=None, **kwargs
+        cls, mu=0.0, sigma=1.0, *, init=None, steps=None, size=None, **kwargs
     ) -> at.TensorVariable:
 
         mu = at.as_tensor_variable(floatX(mu))
         sigma = at.as_tensor_variable(floatX(sigma))
 
-        steps = get_steps_from_shape(steps, kwargs.get("shape", None), step_shape_offset=1)
+        steps = get_steps(
+            steps=steps,
+            shape=kwargs.get("shape", None),
+            step_shape_offset=1,
+        )
+        if steps is None:
+            raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps))
 
         # If no scalar distribution is passed then initialize with a Normal of same mu and sigma
@@ -245,6 +274,7 @@ def dist(
                 and init.owner.op.ndim_supp == 0
             ):
                 raise TypeError("init must be a univariate distribution variable")
+            check_dist_not_registered(init)
 
         # Ignores logprob of init var because that's accounted for in the logp method
         init = ignore_logprob(init)
@@ -340,6 +370,9 @@ class AR(SymbolicDistribution):
     ar_order: int, optional
         Order of the AR process. Inferred from length of the last dimension of rho, if
         possible. ar_order = rho.shape[-1] if constant else rho.shape[-1] - 1
+    steps : int, optional
+        Number of steps in AR process (steps > 0). Only needed if size is used to
+        specify distribution
 
     Notes
     -----
@@ -360,6 +393,20 @@ class AR(SymbolicDistribution):
 
     """
 
+    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,
+            shape=None,  # Shape will be checked in `cls.dist`
+            dims=kwargs.get("dims", None),
+            observed=kwargs.get("observed", None),
+            step_shape_offset=ar_order,
+        )
+        return super().__new__(
+            cls, name, rhos, *args, steps=steps, constant=constant, ar_order=ar_order, **kwargs
+        )
+
     @classmethod
     def dist(
         cls,
@@ -384,32 +431,12 @@ def dist(
             )
             init_dist = kwargs["init"]
 
-        steps = get_steps_from_shape(steps, kwargs.get("shape", None))
+        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)
+        if steps is None:
+            raise ValueError("Must specify steps or shape parameter")
         steps = at.as_tensor_variable(intX(steps), ndim=0)
 
-        if ar_order is None:
-            # If ar_order is not specified we do constant folding on the shape of rhos
-            # to retrieve it. For example, this will detect that
-            # Normal(size=(5, 3)).shape[-1] == 3, which is not known by Aesara before.
-            shape_fg = FunctionGraph(
-                outputs=[rhos.shape[-1]],
-                features=[ShapeFeature()],
-                clone=True,
-            )
-            (folded_shape,) = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs
-            folded_shape = getattr(folded_shape, "data", None)
-            if folded_shape is None:
-                raise ValueError(
-                    "Could not infer ar_order from last dimension of rho. Pass it "
-                    "explictily or make sure rho have a static shape"
-                )
-            ar_order = int(folded_shape) - int(constant)
-            if ar_order < 1:
-                raise ValueError(
-                    "Inferred ar_order is smaller than 1. Increase the last dimension "
-                    "of rho or remove constant_term"
-                )
-
         if init_dist is not None:
             if not isinstance(init_dist, TensorVariable) or not isinstance(
                 init_dist.owner.op, RandomVariable
@@ -433,6 +460,41 @@ def dist(
 
         return super().dist([rhos, sigma, init_dist, steps, ar_order, constant], **kwargs)
 
+    @classmethod
+    def _get_ar_order(cls, rhos: TensorVariable, ar_order: Optional[int], constant: bool) -> int:
+        """Compute ar_order given inputs
+
+        If ar_order is not specified we do constant folding on the shape of rhos
+        to retrieve it. For example, this will detect that
+        Normal(size=(5, 3)).shape[-1] == 3, which is not known by Aesara before.
+
+        Raises
+        ------
+        ValueError
+            If inferred ar_order cannot be inferred from rhos or if it is less than 1
+        """
+        if ar_order is None:
+            shape_fg = FunctionGraph(
+                outputs=[rhos.shape[-1]],
+                features=[ShapeFeature()],
+                clone=True,
+            )
+            (folded_shape,) = optimize_graph(shape_fg, custom_opt=topo_constant_folding).outputs
+            folded_shape = getattr(folded_shape, "data", None)
+            if folded_shape is None:
+                raise ValueError(
+                    "Could not infer ar_order from last dimension of rho. Pass it "
+                    "explictily or make sure rho have a static shape"
+                )
+            ar_order = int(folded_shape) - int(constant)
+            if ar_order < 1:
+                raise ValueError(
+                    "Inferred ar_order is smaller than 1. Increase the last dimension "
+                    "of rho or remove constant_term"
+                )
+
+        return ar_order
+
     @classmethod
     def num_rngs(cls, *args, **kwargs):
         return 2
@@ -496,7 +558,7 @@ def step(*args):
             fn=step,
             outputs_info=[{"initial": init_.T, "taps": range(-ar_order, 0)}],
             non_sequences=[rhos_bcast_.T[::-1], sigma_.T, noise_rng],
-            n_steps=at.max((0, steps_ - ar_order)),
+            n_steps=steps_,
             strict=True,
         )
         (noise_next_rng,) = tuple(innov_updates_.values())

pymc/model.py

@@ -50,10 +50,10 @@
 
 from pymc.aesaraf import (
     compile_pymc,
+    convert_observed_data,
     gradient,
     hessian,
     inputvars,
-    pandas_to_array,
     rvs_to_value_vars,
 )
 from pymc.blocking import DictToArrayBijection, RaveledVars
@@ -1158,7 +1158,7 @@ def set_data(
 
         if isinstance(values, list):
             values = np.array(values)
-        values = pandas_to_array(values)
+        values = convert_observed_data(values)
         dims = self.RV_dims.get(name, None) or ()
         coords = coords or {}
 
@@ -1290,7 +1290,7 @@ def make_obs_var(
 
         """
         name = rv_var.name
-        data = pandas_to_array(data).astype(rv_var.dtype)
+        data = convert_observed_data(data).astype(rv_var.dtype)
 
         if data.ndim != rv_var.ndim:
             raise ShapeError(