Create dispatched resize_dist

Author: ricardoV94

pymc/aesaraf.py

@@ -31,7 +31,7 @@
 import scipy.sparse as sps
 
 from aeppl.logprob import CheckParameterValue
-from aesara import config, scalar
+from aesara import scalar
 from aesara.compile.mode import Mode, get_mode
 from aesara.gradient import grad
 from aesara.graph import local_optimizer
@@ -45,17 +45,15 @@
     walk,
 )
 from aesara.graph.fg import FunctionGraph
-from aesara.graph.op import Op, compute_test_value
+from aesara.graph.op import Op
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.scalar.basic import Cast
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
-from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorConstant, TensorVariable
 
-from pymc.exceptions import ShapeError
 from pymc.vartypes import continuous_types, int_types, isgenerator, typefilter
 
 PotentialShapeType = Union[
@@ -142,64 +140,6 @@ def pandas_to_array(data):
         return floatX(ret)
 
 
-def change_rv_size(
-    rv: TensorVariable,
-    new_size: PotentialShapeType,
-    expand: Optional[bool] = False,
-) -> TensorVariable:
-    """Change or expand the size of a `RandomVariable`.
-
-    Parameters
-    ==========
-    rv
-        The old `RandomVariable` output.
-    new_size
-        The new size.
-    expand:
-        Expand the existing size by `new_size`.
-
-    """
-    # Check the dimensionality of the `new_size` kwarg
-    new_size_ndim = np.ndim(new_size)
-    if new_size_ndim > 1:
-        raise ShapeError("The `new_size` must be ≤1-dimensional.", actual=new_size_ndim)
-    elif new_size_ndim == 0:
-        new_size = (new_size,)
-
-    # Extract the RV node that is to be resized, together with its inputs, name and tag
-    if isinstance(rv.owner.op, SpecifyShape):
-        rv = rv.owner.inputs[0]
-    rv_node = rv.owner
-    rng, size, dtype, *dist_params = rv_node.inputs
-    name = rv.name
-    tag = rv.tag
-
-    if expand:
-        shape = tuple(rv_node.op._infer_shape(size, dist_params))
-        size = shape[: len(shape) - rv_node.op.ndim_supp]
-        new_size = tuple(new_size) + tuple(size)
-
-    # Make sure the new size is a tensor. This dtype-aware conversion helps
-    # to not unnecessarily pick up a `Cast` in some cases (see #4652).
-    new_size = at.as_tensor(new_size, ndim=1, dtype="int64")
-
-    new_rv_node = rv_node.op.make_node(rng, new_size, dtype, *dist_params)
-    new_rv = new_rv_node.outputs[-1]
-    new_rv.name = name
-    for k, v in tag.__dict__.items():
-        new_rv.tag.__dict__.setdefault(k, v)
-
-    # Update "traditional" rng default_update, if that was set for old RV
-    default_update = getattr(rng, "default_update", None)
-    if default_update is not None and default_update is rv_node.outputs[0]:
-        rng.default_update = new_rv_node.outputs[0]
-
-    if config.compute_test_value != "off":
-        compute_test_value(new_rv_node)
-
-    return new_rv
-
-
 def extract_rv_and_value_vars(
     var: TensorVariable,
 ) -> Tuple[TensorVariable, TensorVariable]:

pymc/distributions/censored.py

@@ -19,7 +19,12 @@
 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.distributions.shape_utils import (
+    _ndim_supp_dist,
+    _resize_dist,
+    ndim_supp_dist,
+    resize_dist,
+)
 from pymc.util import check_dist_not_registered
 
 
@@ -90,22 +95,12 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None):
         rv_out.tag.upper = upper
 
         if size is not None:
-            rv_out = cls.change_size(rv_out, size)
+            rv_out = resize_dist(rv_out, size)
         if rngs is not None:
             rv_out = cls.change_rngs(rv_out, rngs)
 
         return rv_out
 
-    @classmethod
-    def change_size(cls, rv, new_size, expand=False):
-        dist_node = rv.tag.dist.owner
-        lower = rv.tag.lower
-        upper = rv.tag.upper
-        rng, old_size, dtype, *dist_params = dist_node.inputs
-        new_size = new_size if not expand else tuple(new_size) + tuple(old_size)
-        new_dist = dist_node.op.make_node(rng, new_size, dtype, *dist_params).default_output()
-        return cls.rv_op(new_dist, lower, upper)
-
     @classmethod
     def change_rngs(cls, rv, new_rngs):
         (new_rng,) = new_rngs
@@ -127,6 +122,15 @@ def ndim_supp_censored(op, dist):
     return 0
 
 
+@_resize_dist.register(Clip)
+def resize_censored(op, rv, new_size, expand=False):
+    dist = rv.tag.dist
+    lower = rv.tag.lower
+    upper = rv.tag.upper
+    new_dist = resize_dist(dist, new_size, expand=expand)
+    return Censored.rv_op(new_dist, lower, upper)
+
+
 @_moment.register(Clip)
 def moment_censored(op, rv, dist, lower, upper):
     moment = at.switch(

pymc/distributions/distribution.py

@@ -33,7 +33,6 @@
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
-from pymc.aesaraf import change_rv_size
 from pymc.distributions.shape_utils import (
     Dims,
     Shape,
@@ -45,6 +44,7 @@
     convert_size,
     find_size,
     ndim_supp_dist,
+    resize_dist,
     resize_from_dims,
     resize_from_observed,
 )
@@ -270,7 +270,7 @@ def __new__(
 
         if resize_shape:
             # A batch size was specified through `dims`, or implied by `observed`.
-            rv_out = change_rv_size(rv=rv_out, new_size=resize_shape, expand=True)
+            rv_out = resize_dist(dist=rv_out, new_size=resize_shape, expand=True)
 
         rv_out = model.register_rv(
             rv_out,
@@ -356,7 +356,7 @@ def dist(
         # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
         if shape is not None and Ellipsis in shape:
             replicate_shape = cast(StrongShape, shape[:-1])
-            rv_out = change_rv_size(rv=rv_out, new_size=replicate_shape, expand=True)
+            rv_out = resize_dist(dist=rv_out, new_size=replicate_shape, expand=True)
 
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
         rv_out.logcdf = _make_nice_attr_error("rv.logcdf(x)", "pm.logcdf(rv, x)")
@@ -400,9 +400,6 @@ 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.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.
 
@@ -413,6 +410,9 @@ def __new__(
             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`
+         _resize_dist
+            Returns an equivalent symbolic distribution with a different size. This is
+            called by `pymc.distrributions.shape_utils.resize_dist`.
 
         Parameters
         ----------

pymc/distributions/mixture.py

@@ -25,13 +25,18 @@
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
 
-from pymc.aesaraf import change_rv_size
 from pymc.distributions import transforms
 from pymc.distributions.continuous import 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 logcdf, logp
-from pymc.distributions.shape_utils import _ndim_supp_dist, ndim_supp_dist, to_tuple
+from pymc.distributions.shape_utils import (
+    _ndim_supp_dist,
+    _resize_dist,
+    ndim_supp_dist,
+    resize_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
@@ -323,27 +328,7 @@ def _resize_components(cls, size, *components):
             mix_size = components[0].shape[mix_axis]
             size = tuple(size) + (mix_size,)
 
-        return [change_rv_size(component, size) for component in components]
-
-    @classmethod
-    def change_size(cls, rv, new_size, expand=False):
-        mix_indexes_rng, weights, *components = rv.owner.inputs
-        rngs = [component.owner.inputs[0] for component in components] + [mix_indexes_rng]
-
-        if expand:
-            component = components[0]
-            # Old size is equal to `shape[:-ndim_supp]`, with care needed for `ndim_supp == 0`
-            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`
-                size_dims -= 1
-            old_size = components[0].shape[:size_dims]
-            new_size = to_tuple(new_size) + tuple(old_size)
-
-        components = cls._resize_components(new_size, *components)
-
-        return cls.rv_op(weights, *components, rngs=rngs, size=None)
+        return [resize_dist(component, size) for component in components]
 
     @classmethod
     def graph_rvs(cls, rv):
@@ -361,6 +346,27 @@ def ndim_supp_marginal_mixture(op, rv):
     return ndim_supp_dist(components[0])
 
 
+@_resize_dist.register(MarginalMixtureRV)
+def resize_marginal_mixture(op, rv, new_size, expand=False):
+    mix_indexes_rng, weights, *components = rv.owner.inputs
+    rngs = [component.owner.inputs[0] for component in components] + [mix_indexes_rng]
+
+    if expand:
+        component = components[0]
+        # Old size is equal to `shape[:-ndim_supp]`, with care needed for `ndim_supp == 0`
+        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`
+            size_dims -= 1
+        old_size = components[0].shape[:size_dims]
+        new_size = to_tuple(new_size) + tuple(old_size)
+
+    components = Mixture._resize_components(new_size, *components)
+
+    return Mixture.rv_op(weights, *components, rngs=rngs, size=None)
+
+
 @_get_measurable_outputs.register(MarginalMixtureRV)
 def _get_measurable_outputs_MarginalMixtureRV(op, node):
     # This tells Aeppl that the second output is the measurable one

pymc/distributions/multivariate.py

@@ -41,7 +41,7 @@
 
 import pymc as pm
 
-from pymc.aesaraf import change_rv_size, floatX, intX
+from pymc.aesaraf import floatX, intX
 from pymc.distributions import transforms
 from pymc.distributions.continuous import (
     BoundedContinuous,
@@ -60,6 +60,7 @@
 from pymc.distributions.shape_utils import (
     broadcast_dist_samples_to,
     ndim_supp_dist,
+    resize_dist,
     rv_size_is_none,
     to_tuple,
 )
@@ -1199,10 +1200,10 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs):
         # Since `eta` and `n` are forced to be scalars we don't need to worry about
         # implied batched dimensions for the time being.
         if ndim_supp_dist(sd_dist) == 0:
-            sd_dist = change_rv_size(sd_dist, to_tuple(size) + (n,))
+            sd_dist = resize_dist(sd_dist, to_tuple(size) + (n,))
         else:
             # The support shape must be `n` but we have no way of controlling it
-            sd_dist = change_rv_size(sd_dist, to_tuple(size))
+            sd_dist = resize_dist(sd_dist, to_tuple(size))
 
         # sd_dist is part of the generative graph, but should be completely ignored
         # by the logp graph, since the LKJ logp explicitly includes these terms.