Allow size broadcasting from dims and observed

Author: ricardoV94

pymc/distributions/distribution.py

@@ -33,18 +33,18 @@
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
-from pymc.aesaraf import change_rv_size
+from pymc.aesaraf import change_rv_size, convert_observed_data
 from pymc.distributions.shape_utils import (
     Dims,
     Shape,
     Size,
+    StrongDims,
     StrongShape,
     convert_dims,
     convert_shape,
     convert_size,
     find_size,
-    resize_from_dims,
-    resize_from_observed,
+    shape_from_dims,
 )
 from pymc.printing import str_for_dist, str_for_symbolic_dist
 from pymc.util import UNSET
@@ -152,29 +152,28 @@ def fn(*args, **kwargs):
 def _make_rv_and_resize_shape(
     *,
     cls,
-    dims: Optional[Dims],
+    dims: Optional[StrongDims],
     model,
     observed,
     args,
     **kwargs,
-) -> Tuple[Variable, Optional[Dims], Optional[Union[np.ndarray, Variable]], StrongShape]:
-    """Creates the RV and processes dims or observed to determine a resize shape."""
-    # Create the RV without dims information, because that's not something tracked at the Aesara level.
-    # If necessary we'll later replicate to a different size implied by already known dims.
-    rv_out = cls.dist(*args, **kwargs)
-    ndim_actual = rv_out.ndim
+) -> Tuple[Variable, StrongShape]:
+    """Creates the RV, possibly using dims or observed to determine a resize shape (if needed)."""
     resize_shape = None
+    size_or_shape = kwargs.get("size") or kwargs.get("shape")
+
+    # Create the RV without dims or observed information
+    rv_out = cls.dist(*args, **kwargs)
 
-    # # `dims` are only available with this API, because `.dist()` can be used
-    # # without a modelcontext and dims are not tracked at the Aesara level.
-    dims = convert_dims(dims)
-    dims_can_resize = kwargs.get("shape", None) is None and kwargs.get("size", None) is None
-    if dims is not None:
-        if dims_can_resize:
-            resize_shape, dims = resize_from_dims(dims, ndim_actual, model)
-    elif observed is not None:
-        resize_shape, observed = resize_from_observed(observed, ndim_actual)
-    return rv_out, dims, observed, resize_shape
+    # Preference is given to size or shape, if not provided we use dims and observed
+    # to resize the variable
+    if not size_or_shape:
+        if dims is not None:
+            resize_shape = shape_from_dims(dims, tuple(rv_out.shape), model)
+        elif observed is not None:
+            resize_shape = tuple(observed.shape)
+
+    return rv_out, resize_shape
 
 
 class Distribution(metaclass=DistributionMeta):
@@ -254,15 +253,20 @@ def __new__(
         if not isinstance(name, string_types):
             raise TypeError(f"Name needs to be a string but got: {name}")
 
-        # Create the RV and process dims and observed to determine
-        # a shape by which the created RV may need to be resized.
-        rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
+        dims = convert_dims(dims)
+        if observed is not None:
+            observed = convert_observed_data(observed)
+
+        # Create the RV, possibly taking into consideration dims and observed to
+        # determine its shape
+        rv_out, resize_shape = _make_rv_and_resize_shape(
             cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs
         )
 
+        # A shape was specified only through `dims`, or implied by `observed`.
         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)
+            resize_size = find_size(shape=resize_shape, size=None, ndim_supp=cls.rv_op.ndim_supp)
+            rv_out = change_rv_size(rv=rv_out, new_size=resize_size, expand=False)
 
         rv_out = model.register_rv(
             rv_out,
@@ -336,11 +340,7 @@ def dist(
         shape = convert_shape(shape)
         size = convert_size(size)
 
-        create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
-            shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp
-        )
-        # Create the RV with a `size` right away.
-        # This is not necessarily the final result.
+        create_size = find_size(shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp)
         rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
 
         rv_out.logp = _make_nice_attr_error("rv.logp(x)", "pm.logp(rv, x)")
@@ -448,19 +448,20 @@ def __new__(
         if not isinstance(name, string_types):
             raise TypeError(f"Name needs to be a string but got: {name}")
 
-        # Create the RV and process dims and observed to determine
-        # a shape by which the created RV may need to be resized.
-        rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
+        dims = convert_dims(dims)
+        if observed is not None:
+            observed = convert_observed_data(observed)
+
+        # Create the RV, possibly taking into consideration dims and observed to
+        # determine its shape
+        rv_out, resize_shape = _make_rv_and_resize_shape(
             cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs
         )
 
+        # A shape was specified only through `dims`, or implied by `observed`.
         if resize_shape:
-            # A batch size was specified through `dims`, or implied by `observed`.
-            rv_out = cls.change_size(
-                rv=rv_out,
-                new_size=resize_shape,
-                expand=True,
-            )
+            resize_size = find_size(shape=resize_shape, size=None, ndim_supp=rv_out.tag.ndim_supp)
+            rv_out = cls.change_size(rv=rv_out, new_size=resize_size, expand=False)
 
         rv_out = model.register_rv(
             rv_out,
@@ -529,18 +530,17 @@ def dist(
         shape = convert_shape(shape)
         size = convert_size(size)
 
-        create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
-            shape=shape, size=size, ndim_supp=cls.ndim_supp(*dist_params)
-        )
-        # Create the RV with a `size` right away.
-        # This is not necessarily the final result.
-        graph = cls.rv_op(*dist_params, size=create_size, **kwargs)
+        ndim_supp = cls.ndim_supp(*dist_params)
+        create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
+        rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
+        # This is needed for resizing from dims in `__new__`
+        rv_out.tag.ndim_supp = ndim_supp
 
         # TODO: Create new attr error stating that these are not available for DerivedDistribution
         # 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)")
         # rv_out.random = _make_nice_attr_error("rv.random()", "rv.eval()")
-        return graph
+        return rv_out
 
 
 @singledispatch

pymc/distributions/shape_utils.py

@@ -18,16 +18,14 @@
 samples from probability distributions for stochastic nodes in PyMC.
 """
 
-from typing import Optional, Sequence, Tuple, Union, cast
+from typing import Optional, Sequence, Tuple, Union
 
 import numpy as np
 
 from aesara.graph.basic import Variable
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
-from pymc.aesaraf import convert_observed_data
-
 __all__ = [
     "to_tuple",
     "shapes_broadcasting",
@@ -471,74 +469,46 @@ def convert_size(size: Size) -> Optional[StrongSize]:
     return size
 
 
-def resize_from_dims(dims: Dims, ndim_implied: int, model) -> Tuple[StrongSize, StrongDims]:
-    """Determines a potential resize shape from a `dims` tuple.
+def shape_from_dims(
+    dims: StrongDims, shape_implied: Sequence[TensorVariable], model
+) -> StrongShape:
+    """Determines shape from a `dims` tuple.
 
     Parameters
     ----------
     dims : array-like
         A vector of dimension names or None.
-    ndim_implied : int
-        Number of RV dimensions that were implied from its inputs alone.
+    shape_implied : tensor_like of int
+        Shape of RV implied from its inputs alone.
     model : pm.Model
         The current model on stack.
 
     Returns
     -------
-    resize_shape : array-like
-        Shape of new dimensions that should be prepended.
     dims : tuple of (str or None)
-        Names or None for all dimensions after resizing.
+        Names or None for all RV dimensions.
     """
-    sdims = cast(StrongDims, dims)
-
-    ndim_resize = len(sdims) - ndim_implied
+    ndim_resize = len(dims) - len(shape_implied)
 
-    # All resize dims must be known already (numerically or symbolically).
-    unknowndim_resize_dims = set(sdims[:ndim_resize]) - set(model.dim_lengths)
+    # Dims must be known already or be inferrable from implied dimensions of the RV
+    unknowndim_resize_dims = set(dims[:ndim_resize]) - set(model.dim_lengths)
     if unknowndim_resize_dims:
         raise KeyError(
             f"Dimensions {unknowndim_resize_dims} are unknown to the model and cannot be used to specify a `size`."
         )
 
     # The numeric/symbolic resize tuple can be created using model.RV_dim_lengths
-    resize_shape: Tuple[Variable, ...] = tuple(
-        model.dim_lengths[dname] for dname in sdims[:ndim_resize]
+    return tuple(
+        model.dim_lengths[dname] if dname in model.dim_lengths else shape_implied[i]
+        for i, dname in enumerate(dims)
     )
-    return resize_shape, sdims
-
-
-def resize_from_observed(
-    observed, ndim_implied: int
-) -> Tuple[StrongSize, Union[np.ndarray, Variable]]:
-    """Determines a potential resize shape from observations.
-
-    Parameters
-    ----------
-    observed : scalar, array-like
-        The value of the `observed` kwarg to the RV creation.
-    ndim_implied : int
-        Number of RV dimensions that were implied from its inputs alone.
-
-    Returns
-    -------
-    resize_shape : array-like
-        Shape of new dimensions that should be prepended.
-    observed : scalar, array-like
-        Observations as numpy array or `Variable`.
-    """
-    if not hasattr(observed, "shape"):
-        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
 
 
 def find_size(
     shape: Optional[StrongShape],
     size: Optional[StrongSize],
     ndim_supp: int,
-) -> Tuple[Optional[StrongSize], Optional[int], Optional[int], int]:
+) -> Optional[StrongSize]:
     """Determines the size keyword argument for creating a Distribution.
 
     Parameters
@@ -553,30 +523,19 @@ def find_size(
 
     Returns
     -------
-    create_size : int, optional
-        The size argument to be passed to the distribution
-    ndim_expected : int, optional
-        Number of dimensions expected after distribution was created
-    ndim_batch : int, optional
-        Number of batch dimensions
-    ndim_supp : int
-        Number of support dimensions
+    size : tuble of int or TensorVariable, optional
+        The size argument for creating the Distribution
     """
 
-    ndim_expected: Optional[int] = None
-    ndim_batch: Optional[int] = None
-    create_size: Optional[StrongSize] = None
+    if size is not None:
+        return size
 
     if shape is not None:
         ndim_expected = len(tuple(shape))
         ndim_batch = ndim_expected - ndim_supp
-        create_size = tuple(shape)[:ndim_batch]
-    elif size is not None:
-        ndim_expected = ndim_supp + len(tuple(size))
-        ndim_batch = ndim_expected - ndim_supp
-        create_size = size
+        return tuple(shape)[:ndim_batch]
 
-    return create_size, ndim_expected, ndim_batch, ndim_supp
+    return None
 
 
 def rv_size_is_none(size: Variable) -> bool:

pymc/model.py

@@ -43,6 +43,8 @@
 from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph.basic import Constant, Variable, graph_inputs
 from aesara.graph.fg import FunctionGraph
+from aesara.scalar import Cast
+from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.rewriting import local_subtensor_rv_lift
 from aesara.tensor.sharedvar import ScalarSharedVariable
 from aesara.tensor.var import TensorConstant, TensorVariable
@@ -1367,6 +1369,12 @@ def register_rv(
                 isinstance(data, Variable)
                 and not isinstance(data, (GenTensorVariable, Minibatch))
                 and data.owner is not None
+                # The only Aesara operation we allow on observed data is type casting
+                # Although we could allow for any graph that does not depend on other RVs
+                and not (
+                    isinstance(data.owner.op, Elemwise)
+                    and isinstance(data.owner.op.scalar_op, Cast)
+                )
             ):
                 raise TypeError(
                     "Variables that depend on other nodes cannot be used for observed data."

pymc/tests/test_data_container.py

@@ -45,7 +45,7 @@ def test_sample(self):
         with pm.Model():
             x_shared = pm.MutableData("x_shared", x)
             b = pm.Normal("b", 0.0, 10.0)
-            pm.Normal("obs", b * x_shared, np.sqrt(1e-2), observed=y)
+            pm.Normal("obs", b * x_shared, np.sqrt(1e-2), observed=y, shape=x_shared.shape)
 
             prior_trace0 = pm.sample_prior_predictive(1000)
             idata = pm.sample(1000, tune=1000, chains=1)

pymc/tests/test_distributions_random.py

@@ -1905,7 +1905,6 @@ def test_density_dist_with_random(self, size):
                 mu,
                 random=lambda mu, rng=None, size=None: rng.normal(loc=mu, scale=1, size=size),
                 observed=np.random.randn(100, *size),
-                size=size,
             )
 
         assert obs.eval().shape == (100,) + size
@@ -1937,7 +1936,6 @@ def test_density_dist_with_random_multivariate(self, size):
                     mean=mu, cov=np.eye(len(mu)), size=size
                 ),
                 observed=np.random.randn(100, *size, supp_shape),
-                size=size,
                 ndims_params=[1],
                 ndim_supp=1,
             )

pymc/tests/test_sampling.py

@@ -1606,7 +1606,7 @@ def test_linear_model(self):
             beta = pm.Normal("beta", 0, 0.1)
             mu = pm.Deterministic("mu", alpha + beta * x)
             sigma = pm.HalfNormal("sigma", 0.1)
-            obs = pm.Normal("obs", mu, sigma, observed=y)
+            obs = pm.Normal("obs", mu, sigma, observed=y, shape=x.shape)
 
         f = compile_forward_sampling_function(
             [obs],
@@ -1624,7 +1624,7 @@ def test_linear_model(self):
             beta = pm.Normal("beta", 0, 0.1)
             mu = pm.Deterministic("mu", alpha + beta * x)
             sigma = pm.HalfNormal("sigma", 0.1)
-            obs = pm.Normal("obs", mu, sigma, observed=y)
+            obs = pm.Normal("obs", mu, sigma, observed=y, shape=x.shape)
 
         f = compile_forward_sampling_function(
             [obs],
@@ -1644,7 +1644,7 @@ def test_nested_observed_model(self):
             beta = pm.Normal("beta", 0, 0.1, size=p.shape)
             mu = pm.Deterministic("mu", beta[category])
             sigma = pm.HalfNormal("sigma", 0.1)
-            pm.Normal("obs", mu, sigma, observed=y)
+            pm.Normal("obs", mu, sigma, observed=y, shape=mu.shape)
 
         f = compile_forward_sampling_function(
             outputs=model.observed_RVs,
@@ -1675,7 +1675,7 @@ def test_volatile_parameters(self):
             mu = pm.Normal("mu", 0, 1)
             nested_mu = pm.Normal("nested_mu", mu, 1, size=10)
             sigma = pm.HalfNormal("sigma", 1)
-            pm.Normal("obs", nested_mu, sigma, observed=y)
+            pm.Normal("obs", nested_mu, sigma, observed=y, shape=nested_mu.shape)
 
         f = compile_forward_sampling_function(
             outputs=model.observed_RVs,