Create dispatched ndim_supp_dist

Author: ricardoV94

pymc/distributions/censored.py

@@ -19,6 +19,7 @@
 from aesara.tensor.random.op import RandomVariable
 
 from pymc.distributions.distribution import SymbolicDistribution, _moment
+from pymc.distributions.shape_utils import _ndim_supp_dist, ndim_supp_dist
 from pymc.util import check_dist_not_registered
 
 
@@ -65,7 +66,7 @@ def dist(cls, dist, lower, upper, **kwargs):
             raise ValueError(
                 f"Censoring dist must be a distribution created via the `.dist()` API, got {type(dist)}"
             )
-        if dist.owner.op.ndim_supp > 0:
+        if ndim_supp_dist(dist) > 0:
             raise NotImplementedError(
                 "Censoring of multivariate distributions has not been implemented yet"
             )
@@ -95,10 +96,6 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
 
         return rv_out
 
-    @classmethod
-    def ndim_supp(cls, *dist_params):
-        return 0
-
     @classmethod
     def change_size(cls, rv, new_size, expand=False):
         dist_node = rv.tag.dist.owner
@@ -124,6 +121,12 @@ def graph_rvs(cls, rv):
         return (rv.tag.dist,)
 
 
+@_ndim_supp_dist.register(Clip)
+def ndim_supp_censored(op, dist):
+    # We only support Censoring of univariate distributions
+    return 0
+
+
 @_moment.register(Clip)
 def moment_censored(op, rv, dist, lower, upper):
     moment = at.switch(

pymc/distributions/distribution.py

@@ -44,6 +44,7 @@
     convert_shape,
     convert_size,
     find_size,
+    ndim_supp_dist,
     resize_from_dims,
     resize_from_observed,
 )
@@ -399,16 +400,20 @@ def __new__(
         cls.rv_op
             Returns a TensorVariable that represents the symbolic distribution
             parametrized by a default set of parameters and a size and rngs arguments
-        cls.ndim_supp
-            Returns the support of the symbolic distribution, given the default
-            parameters. This may not always be constant, for instance if the symbolic
-            distribution can be defined based on an arbitrary base distribution.
         cls.change_size
             Returns an equivalent symbolic distribution with a different size. This is
             analogous to `pymc.aesaraf.change_rv_size` for `RandomVariable`s.
         cls.graph_rvs
             Returns base RVs in a symbolic distribution.
 
+        Furthermore, Censored distributions must have a dispatch version of the following
+        functions for correct behavior in PyMC:
+        _ndim_supp_dist
+            Returns the support of the symbolic distribution. This may not always be
+            constant, for instance if the symbolic distribution can be defined based
+            on an arbitrary base distribution. This is called by
+            `pymc.distributions.shape_utils.ndim_supp_dist`
+
         Parameters
         ----------
         cls : type
@@ -559,8 +564,11 @@ def dist(
         shape = convert_shape(shape)
         size = convert_size(size)
 
+        # Create a temporary dist to obtain the ndim_supp
+        ndim_supp = ndim_supp_dist(cls.rv_op(*dist_params, size=size))
+
         create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
-            shape=shape, size=size, ndim_supp=cls.ndim_supp(*dist_params)
+            shape=shape, size=size, ndim_supp=ndim_supp
         )
         # Create the RV with a `size` right away.
         # This is not necessarily the final result.

pymc/distributions/mixture.py

@@ -31,7 +31,7 @@
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import SymbolicDistribution, _moment, moment
 from pymc.distributions.logprob import logcdf, logp
-from pymc.distributions.shape_utils import to_tuple
+from pymc.distributions.shape_utils import _ndim_supp_dist, ndim_supp_dist, to_tuple
 from pymc.distributions.transforms import _default_transform
 from pymc.util import check_dist_not_registered
 from pymc.vartypes import continuous_types, discrete_types
@@ -188,7 +188,7 @@ def dist(cls, w, comp_dists, **kwargs):
                     f"Component dist must be a distribution created via the `.dist()` API, got {type(dist)}"
                 )
             check_dist_not_registered(dist)
-            components_ndim_supp.add(dist.owner.op.ndim_supp)
+            components_ndim_supp.add(ndim_supp_dist(dist))
 
         if len(components_ndim_supp) > 1:
             raise ValueError(
@@ -209,7 +209,7 @@ def rv_op(cls, weights, *components, size=None, rngs=None):
             mix_indexes_rng = aesara.shared(np.random.default_rng())
 
         single_component = len(components) == 1
-        ndim_supp = components[0].owner.op.ndim_supp
+        ndim_supp = ndim_supp_dist(components[0])
 
         if size is not None:
             components = cls._resize_components(size, *components)
@@ -319,17 +319,12 @@ def _resize_components(cls, size, *components):
         if len(components) == 1:
             # If we have a single component, we need to keep the length of the mixture
             # axis intact, because that's what determines the number of mixture components
-            mix_axis = -components[0].owner.op.ndim_supp - 1
+            mix_axis = -ndim_supp_dist(components[0]) - 1
             mix_size = components[0].shape[mix_axis]
             size = tuple(size) + (mix_size,)
 
         return [change_rv_size(component, size) for component in components]
 
-    @classmethod
-    def ndim_supp(cls, weights, *components):
-        # We already checked that all components have the same support dimensionality
-        return components[0].owner.op.ndim_supp
-
     @classmethod
     def change_size(cls, rv, new_size, expand=False):
         mix_indexes_rng, weights, *components = rv.owner.inputs
@@ -338,7 +333,7 @@ def change_size(cls, rv, new_size, expand=False):
         if expand:
             component = components[0]
             # Old size is equal to `shape[:-ndim_supp]`, with care needed for `ndim_supp == 0`
-            size_dims = component.ndim - component.owner.op.ndim_supp
+            size_dims = component.ndim - ndim_supp_dist(component)
             if len(components) == 1:
                 # If we have a single component, new size should ignore the mixture axis
                 # dimension, as that is not touched by `_resize_components`
@@ -359,6 +354,13 @@ def graph_rvs(cls, rv):
         return (*rv.owner.inputs[2:], rv)
 
 
+@_ndim_supp_dist.register(MarginalMixtureRV)
+def ndim_supp_marginal_mixture(op, rv):
+    # We already checked that all components have the same support dimensionality
+    components = rv.owner.inputs[2:]
+    return ndim_supp_dist(components[0])
+
+
 @_get_measurable_outputs.register(MarginalMixtureRV)
 def _get_measurable_outputs_MarginalMixtureRV(op, node):
     # This tells Aeppl that the second output is the measurable one
@@ -372,7 +374,7 @@ def marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs):
     # single component
     if len(components) == 1:
         # Need to broadcast value across mixture axis
-        mix_axis = -components[0].owner.op.ndim_supp - 1
+        mix_axis = -ndim_supp_dist(components[0]) - 1
         components_logp = logp(components[0], at.expand_dims(value, mix_axis))
     else:
         components_logp = at.stack(
@@ -405,7 +407,7 @@ def marginal_mixture_logcdf(op, value, rng, weights, *components, **kwargs):
     # single component
     if len(components) == 1:
         # Need to broadcast value across mixture axis
-        mix_axis = -components[0].owner.op.ndim_supp - 1
+        mix_axis = -ndim_supp_dist(components[0]) - 1
         components_logcdf = logcdf(components[0], at.expand_dims(value, mix_axis))
     else:
         components_logcdf = at.stack(
@@ -434,7 +436,7 @@ def marginal_mixture_logcdf(op, value, rng, weights, *components, **kwargs):
 
 @_moment.register(MarginalMixtureRV)
 def marginal_mixture_moment(op, rv, rng, weights, *components):
-    ndim_supp = components[0].owner.op.ndim_supp
+    ndim_supp = ndim_supp_dist(components[0])
     weights = at.shape_padright(weights, ndim_supp)
     mix_axis = -ndim_supp - 1
 

pymc/distributions/multivariate.py

@@ -59,6 +59,7 @@
 from pymc.distributions.distribution import Continuous, Discrete, moment
 from pymc.distributions.shape_utils import (
     broadcast_dist_samples_to,
+    ndim_supp_dist,
     rv_size_is_none,
     to_tuple,
 )
@@ -1187,7 +1188,7 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs):
             isinstance(sd_dist, Variable)
             and sd_dist.owner is not None
             and isinstance(sd_dist.owner.op, RandomVariable)
-            and sd_dist.owner.op.ndim_supp < 2
+            and ndim_supp_dist(sd_dist) < 2
         ):
             raise TypeError("sd_dist must be a scalar or vector distribution variable")
 
@@ -1197,7 +1198,7 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs):
         # diagonal element.
         # Since `eta` and `n` are forced to be scalars we don't need to worry about
         # implied batched dimensions for the time being.
-        if sd_dist.owner.op.ndim_supp == 0:
+        if ndim_supp_dist(sd_dist) == 0:
             sd_dist = change_rv_size(sd_dist, to_tuple(size) + (n,))
         else:
             # The support shape must be `n` but we have no way of controlling it

pymc/distributions/shape_utils.py

@@ -17,12 +17,15 @@
 A collection of common shape operations needed for broadcasting
 samples from probability distributions for stochastic nodes in PyMC.
 """
-
+from functools import singledispatch
 from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
 from aesara.graph.basic import Variable
+from aesara.graph.op import Op
+from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
@@ -619,3 +622,23 @@ def find_size(
 def rv_size_is_none(size: Variable) -> bool:
     """Check wether an rv size is None (ie., at.Constant([]))"""
     return size.type.shape == (0,)  # type: ignore [attr-defined]
+
+
+@singledispatch
+def _ndim_supp_dist(op: Op, dist: TensorVariable) -> int:
+    raise TypeError(f"ndim_supp not known for Op {op}")
+
+
+def ndim_supp_dist(dist: TensorVariable) -> int:
+    return _ndim_supp_dist(dist.owner.op, dist)
+
+
+@_ndim_supp_dist.register(RandomVariable)
+def ndim_supp_rv(op: Op, rv: TensorVariable):
+    return op.ndim_supp
+
+
+@_ndim_supp_dist.register(Elemwise)
+def ndim_supp_elemwise(op: Op, *args, **kwargs):
+    """For Elemwise Ops, dispatch on respective scalar_op"""
+    return _ndim_supp_dist(op.scalar_op, *args, **kwargs)

pymc/distributions/timeseries.py

@@ -24,7 +24,7 @@
 from pymc.distributions import distribution, logprob, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.dist_math import check_parameters
-from pymc.distributions.shape_utils import to_tuple
+from pymc.distributions.shape_utils import ndim_supp_dist, to_tuple
 from pymc.util import check_dist_not_registered
 
 __all__ = [
@@ -175,7 +175,7 @@ def dist(
                 isinstance(init, at.TensorVariable)
                 and init.owner is not None
                 and isinstance(init.owner.op, RandomVariable)
-                and init.owner.op.ndim_supp == 0
+                and ndim_supp_dist(init) == 0
             ):
                 raise TypeError("init must be a univariate distribution variable")
 

pymc/model.py

@@ -43,6 +43,7 @@
 from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph.basic import Constant, Variable, graph_inputs
 from aesara.graph.fg import FunctionGraph
+from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.opt import local_subtensor_rv_lift
 from aesara.tensor.random.var import RandomStateSharedVariable
 from aesara.tensor.sharedvar import ScalarSharedVariable
@@ -1330,6 +1331,12 @@ def make_obs_var(
             )
             warnings.warn(impute_message, ImputationWarning)
 
+            # TODO: Add test for this
+            if not isinstance(rv_var.owner.op, RandomVariable):
+                raise NotImplementedError(
+                    f"Automatic inputation is only supported for RandomVariables, but {rv_var} is of type {rv_var.owner.op}"
+                )
+
             if rv_var.owner.op.ndim_supp > 0:
                 raise NotImplementedError(
                     f"Automatic inputation is only supported for univariate RandomVariables, but {rv_var} is multivariate"