diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index dd5968df7e..d0b6747662 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -19,7 +19,7 @@ from abc import ABCMeta from functools import singledispatch -from typing import Callable, Optional, Sequence, Tuple, Union, cast +from typing import Callable, Optional, Sequence, Tuple, Union import aesara import numpy as np @@ -33,19 +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, - WeakDims, 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 @@ -150,35 +149,33 @@ def fn(*args, **kwargs): return fn -def _make_rv_and_resize_shape( +def _make_rv_and_resize_shape_from_dims( *, cls, - dims: Optional[Dims], + dims: Optional[StrongDims], model, observed, args, **kwargs, -) -> Tuple[Variable, Optional[WeakDims], 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. +) -> Tuple[Variable, StrongShape]: + """Creates the RV, possibly using dims or observed to determine a resize shape (if needed).""" + resize_shape_from_dims = None + size_or_shape = kwargs.get("size") or kwargs.get("shape") + + # Preference is given to size or shape. If not specified, we rely on dims and + # finally, observed, to determine the shape of the variable. Because dims can be + # specified on the fly, we need a two-step process where we first create the RV + # without dims information and then resize it. + if not size_or_shape and observed is not None: + kwargs["shape"] = tuple(observed.shape) + + # Create the RV without dims information rv_out = cls.dist(*args, **kwargs) - ndim_actual = rv_out.ndim - resize_shape = None - - # # `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 Ellipsis in dims: - # Replace ... with None entries to match the actual dimensionality. - dims = (*dims[:-1], *[None] * ndim_actual)[:ndim_actual] - elif observed is not None: - resize_shape, observed = resize_from_observed(observed, ndim_actual) - return rv_out, dims, observed, resize_shape + + if not size_or_shape and dims is not None: + resize_shape_from_dims = shape_from_dims(dims, tuple(rv_out.shape), model) + + return rv_out, resize_shape_from_dims class Distribution(metaclass=DistributionMeta): @@ -212,7 +209,8 @@ def __new__( rng : optional Random number generator to use with the RandomVariable. dims : tuple, optional - A tuple of dimension names known to the model. + A tuple of dimension names known to the model. When shape is not provided, + the shape of dims is used to define the shape of the variable. initval : optional Numeric or symbolic untransformed initial value of matching shape, or one of the following initial value strategies: "moment", "prior". @@ -220,6 +218,8 @@ def __new__( or moment-based initial values in the transformed space. observed : optional Observed data to be passed when registering the random variable in the model. + When neither shape nor dims is provided, the shape of observed is used to + define the shape of the variable. See ``Model.register_rv``. total_size : float, optional See ``Model.register_rv``. @@ -258,15 +258,21 @@ 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, without taking `dims` into consideration + rv_out, resize_shape_from_dims = _make_rv_and_resize_shape_from_dims( cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs ) - 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 variable based on `dims` information + if resize_shape_from_dims: + resize_size_from_dims = find_size( + shape=resize_shape_from_dims, size=None, ndim_supp=cls.rv_op.ndim_supp + ) + rv_out = change_rv_size(rv=rv_out, new_size=resize_size_from_dims, expand=False) rv_out = model.register_rv( rv_out, @@ -286,7 +292,7 @@ def __new__( 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()") + rv_out.random = _make_nice_attr_error("rv.random()", "pm.draw(rv)") return rv_out @classmethod @@ -305,9 +311,6 @@ def dist( The inputs to the `RandomVariable` `Op`. shape : int, tuple, Variable, optional A tuple of sizes for each dimension of the new RV. - - An Ellipsis (...) may be inserted in the last position to short-hand refer to - all the dimensions that the RV would get if no shape/size/dims were passed at all. **kwargs Keyword arguments that will be forwarded to the Aesara RV Op. Most prominently: ``size`` or ``dtype``. @@ -343,21 +346,12 @@ 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) - # 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.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()") + rv_out.random = _make_nice_attr_error("rv.random()", "pm.draw(rv)") return rv_out @@ -414,7 +408,8 @@ def __new__( name : str Name for the new model variable. dims : tuple, optional - A tuple of dimension names known to the model. + A tuple of dimension names known to the model. When shape is not provided, + the shape of dims is used to define the shape of the variable. initval : optional Numeric or symbolic untransformed initial value of matching shape, or one of the following initial value strategies: "moment", "prior". @@ -422,6 +417,8 @@ def __new__( symbolic or moment-based initial values in the transformed space. observed : optional Observed data to be passed when registering the random variable in the model. + When neither shape nor dims is provided, the shape of observed is used to + define the shape of the variable. See ``Model.register_rv``. total_size : float, optional See ``Model.register_rv``. @@ -460,19 +457,21 @@ 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, without taking `dims` into consideration + rv_out, resize_shape_from_dims = _make_rv_and_resize_shape_from_dims( cls=cls, dims=dims, model=model, observed=observed, args=args, **kwargs ) - 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 variable based on `dims` information + if resize_shape_from_dims: + resize_size_from_dims = find_size( + shape=resize_shape_from_dims, size=None, ndim_supp=rv_out.tag.ndim_supp ) + rv_out = cls.change_size(rv=rv_out, new_size=resize_size_from_dims, expand=False) rv_out = model.register_rv( rv_out, @@ -489,6 +488,10 @@ def __new__( functools.partial(str_for_symbolic_dist, formatting="latex"), rv_out ) + 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()", "pm.draw(rv)") + return rv_out @classmethod @@ -508,8 +511,6 @@ def dist( The inputs to the `RandomVariable` `Op`. shape : int, tuple, Variable, optional A tuple of sizes for each dimension of the new RV. - An Ellipsis (...) may be inserted in the last position to short-hand refer to - all the dimensions that the RV would get if no shape/size/dims were passed at all. size : int, tuple, Variable, optional For creating the RV like in Aesara/NumPy. @@ -543,23 +544,16 @@ 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) - - # 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]) - graph = cls.change_size(rv=graph, new_size=replicate_shape, expand=True) - - # 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 + 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 + + 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()", "pm.draw(rv)") + return rv_out @singledispatch diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 23c9237199..21a990c2ac 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -18,7 +18,7 @@ samples from probability distributions for stochastic nodes in PyMC. """ -from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union, cast +from typing import Optional, Sequence, Tuple, Union import numpy as np @@ -26,8 +26,6 @@ from aesara.tensor.var import TensorVariable from typing_extensions import TypeAlias -from pymc.aesaraf import convert_observed_data - __all__ = [ "to_tuple", "shapes_broadcasting", @@ -409,34 +407,18 @@ def broadcast_dist_samples_to(to_shape, samples, size=None): return [np.broadcast_to(o, to_shape) for o in samples] -# Workaround to annotate the Ellipsis type, posted by the BDFL himself. -# See https://github.com/python/typing/issues/684#issuecomment-548203158 -if TYPE_CHECKING: - from enum import Enum - - class ellipsis(Enum): - Ellipsis = "..." - - Ellipsis = ellipsis.Ellipsis -else: - ellipsis = type(Ellipsis) - # User-provided can be lazily specified as scalars -Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable, ellipsis]]] -Dims: TypeAlias = Union[str, Sequence[Optional[Union[str, ellipsis]]]] +Shape: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable]]] +Dims: TypeAlias = Union[str, Sequence[Optional[str]]] Size: TypeAlias = Union[int, TensorVariable, Sequence[Union[int, Variable]]] # After conversion to vectors -WeakShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable, ellipsis], ...]] -WeakDims: TypeAlias = Tuple[Optional[Union[str, ellipsis]], ...] - -# After Ellipsis were substituted StrongShape: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]] StrongDims: TypeAlias = Sequence[Optional[str]] StrongSize: TypeAlias = Union[TensorVariable, Tuple[Union[int, Variable], ...]] -def convert_dims(dims: Optional[Dims]) -> Optional[WeakDims]: +def convert_dims(dims: Optional[Dims]) -> Optional[StrongDims]: """Process a user-provided dims variable into None or a valid dims tuple.""" if dims is None: return None @@ -448,13 +430,10 @@ def convert_dims(dims: Optional[Dims]) -> Optional[WeakDims]: else: raise ValueError(f"The `dims` parameter must be a tuple, str or list. Actual: {type(dims)}") - if any(d == Ellipsis for d in dims[:-1]): - raise ValueError(f"Ellipsis in `dims` may only appear in the last position. Actual: {dims}") - return dims -def convert_shape(shape: Shape) -> Optional[WeakShape]: +def convert_shape(shape: Shape) -> Optional[StrongShape]: """Process a user-provided shape variable into None or a valid shape object.""" if shape is None: return None @@ -468,10 +447,6 @@ def convert_shape(shape: Shape) -> Optional[WeakShape]: raise ValueError( f"The `shape` parameter must be a tuple, TensorVariable, int or list. Actual: {type(shape)}" ) - if isinstance(shape, tuple) and any(s == Ellipsis for s in shape[:-1]): - raise ValueError( - f"Ellipsis in `shape` may only appear in the last position. Actual: {shape}" - ) return shape @@ -490,85 +465,50 @@ def convert_size(size: Size) -> Optional[StrongSize]: raise ValueError( f"The `size` parameter must be a tuple, TensorVariable, int or list. Actual: {type(size)}" ) - if isinstance(size, tuple) and Ellipsis in size: - raise ValueError(f"The `size` parameter cannot contain an Ellipsis. Actual: {size}") return size -def resize_from_dims(dims: WeakDims, 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, None or Ellipsis. - ndim_implied : int - Number of RV dimensions that were implied from its inputs alone. + A vector of dimension names or None. + 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. """ - if Ellipsis in dims: - # Auto-complete the dims tuple to the full length. - # We don't have a way to know the names of implied - # dimensions, so they will be `None`. - dims = (*dims[:-1], *[None] * ndim_implied) - sdims = cast(StrongDims, dims) + ndim_resize = len(dims) - len(shape_implied) - ndim_resize = len(sdims) - ndim_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[WeakShape], + 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 @@ -583,37 +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: - if Ellipsis in shape: - # Ellipsis short-hands all implied dimensions. Therefore - # we don't know how many dimensions to expect. - ndim_expected = ndim_batch = None - # Create the RV with its implied shape and resize later - create_size = None - else: - 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_expected = len(tuple(shape)) 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: diff --git a/pymc/model.py b/pymc/model.py index 7e378908a9..037385d534 100644 --- a/pymc/model.py +++ b/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." diff --git a/pymc/model_graph.py b/pymc/model_graph.py index 8e28aaa971..93f8c3a9af 100644 --- a/pymc/model_graph.py +++ b/pymc/model_graph.py @@ -42,16 +42,25 @@ def get_parent_names(self, var: TensorVariable) -> Set[VarName]: if var.owner is None or var.owner.inputs is None: return set() + def _filter_non_parameter_inputs(var): + node = var.owner + if isinstance(node.op, RandomVariable): + # Filter out rng, dtype and size parameters or RandomVariable nodes + return node.inputs[3:] + else: + # Otherwise return all inputs + return node.inputs + def _expand(x): if x.name: return [x] if isinstance(x.owner, Apply): - return reversed(x.owner.inputs) + return reversed(_filter_non_parameter_inputs(x)) return [] parents = { get_var_name(x) - for x in walk(nodes=var.owner.inputs, expand=_expand) + for x in walk(nodes=_filter_non_parameter_inputs(var), expand=_expand) # Only consider nodes that are in the named model variables. if x.name and x.name in self._all_var_names } diff --git a/pymc/tests/test_data_container.py b/pymc/tests/test_data_container.py index bfa2b191ba..4ccd86426b 100644 --- a/pymc/tests/test_data_container.py +++ b/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) diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index 50b9c3c90d..af137af859 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -3145,7 +3145,7 @@ def test_distinct_rvs(): [ ("logp", r"pm.logp\(rv, x\)"), ("logcdf", r"pm.logcdf\(rv, x\)"), - ("random", r"rv.eval\(\)"), + ("random", r"pm.draw\(rv\)"), ], ) def test_logp_gives_migration_instructions(method, newcode): diff --git a/pymc/tests/test_distributions_random.py b/pymc/tests/test_distributions_random.py index e1aee47f88..56d235f76f 100644 --- a/pymc/tests/test_distributions_random.py +++ b/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, ) diff --git a/pymc/tests/test_sampling.py b/pymc/tests/test_sampling.py index 31bb3dd881..e7a1073dbd 100644 --- a/pymc/tests/test_sampling.py +++ b/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, diff --git a/pymc/tests/test_shape_handling.py b/pymc/tests/test_shape_handling.py index 9714a45492..4ab3b93b4f 100644 --- a/pymc/tests/test_shape_handling.py +++ b/pymc/tests/test_shape_handling.py @@ -204,16 +204,14 @@ def test_broadcast_dist_samples_to(self, samples_to_broadcast_to): class TestShapeDimsSize: - @pytest.mark.parametrize("param_shape", [(), (3,)]) + @pytest.mark.parametrize("param_shape", [(), (2,)]) @pytest.mark.parametrize("batch_shape", [(), (3,)]) @pytest.mark.parametrize( "parametrization", [ "implicit", "shape", - "shape...", "dims", - "dims...", "size", ], ) @@ -249,69 +247,60 @@ def test_param_and_batch_shape_combos( if parametrization == "shape": rv = pm.Normal("rv", mu=mu, shape=batch_shape + param_shape) assert rv.eval().shape == expected_shape - elif parametrization == "shape...": - rv = pm.Normal("rv", mu=mu, shape=(*batch_shape, ...)) - assert rv.eval().shape == batch_shape + param_shape elif parametrization == "dims": rv = pm.Normal("rv", mu=mu, dims=batch_dims + param_dims) assert rv.eval().shape == expected_shape - elif parametrization == "dims...": - rv = pm.Normal("rv", mu=mu, dims=(*batch_dims, ...)) - n_size = len(batch_shape) - n_implied = len(param_shape) - ndim = n_size + n_implied - assert len(pmodel.RV_dims["rv"]) == ndim, pmodel.RV_dims - assert len(pmodel.RV_dims["rv"][:n_size]) == len(batch_dims) - assert len(pmodel.RV_dims["rv"][n_size:]) == len(param_dims) - if n_implied > 0: - assert pmodel.RV_dims["rv"][-1] is None elif parametrization == "size": rv = pm.Normal("rv", mu=mu, size=batch_shape + param_shape) assert rv.eval().shape == expected_shape else: raise NotImplementedError("Invalid test case parametrization.") - @pytest.mark.parametrize("ellipsis_in", ["none", "shape", "dims", "both"]) - def test_simultaneous_shape_and_dims(self, ellipsis_in): + def test_broadcast_by_dims(self): + with pm.Model(coords={"broadcast_dim": range(3)}) as m: + x = pm.Normal("x", mu=np.zeros((1,)), dims=("broadcast_dim",)) + assert x.eval().shape == (3,) + + def test_broadcast_by_observed(self): + with pm.Model() as m: + x = pm.Normal("x", mu=np.zeros((1,)), observed=np.zeros((3,))) + assert x.eval().shape == (3,) + + def test_simultaneous_shape_and_dims(self): with pm.Model() as pmodel: x = pm.ConstantData("x", [1, 2, 3], dims="ddata") - if ellipsis_in == "none": - # The shape and dims tuples correspond to each other. - # Note: No checks are performed that implied shape (x), shape and dims actually match. - y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", "ddata")) - assert pmodel.RV_dims["y"] == ("dshape", "ddata") - elif ellipsis_in == "shape": - y = pm.Normal("y", mu=x, shape=(2, ...), dims=("dshape", "ddata")) - assert pmodel.RV_dims["y"] == ("dshape", "ddata") - elif ellipsis_in == "dims": - y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", ...)) - assert pmodel.RV_dims["y"] == ("dshape", None) - elif ellipsis_in == "both": - y = pm.Normal("y", mu=x, shape=(2, ...), dims=("dshape", ...)) - assert pmodel.RV_dims["y"] == ("dshape", None) + # The shape and dims tuples correspond to each other. + # Note: No checks are performed that implied shape (x), shape and dims actually match. + y = pm.Normal("y", mu=x, shape=(2, 3), dims=("dshape", "ddata")) + assert pmodel.RV_dims["y"] == ("dshape", "ddata") assert "dshape" in pmodel.dim_lengths assert y.eval().shape == (2, 3) - @pytest.mark.parametrize("with_dims_ellipsis", [False, True]) - def test_simultaneous_size_and_dims(self, with_dims_ellipsis): + def test_simultaneous_size_and_dims(self): with pm.Model() as pmodel: x = pm.ConstantData("x", [1, 2, 3], dims="ddata") assert "ddata" in pmodel.dim_lengths # Size does not include support dims, so this test must use a dist with support dims. kwargs = dict(name="y", size=(2, 3), mu=at.ones((3, 4)), cov=at.eye(4)) - if with_dims_ellipsis: - y = pm.MvNormal(**kwargs, dims=("dsize", ...)) - assert pmodel.RV_dims["y"] == ("dsize", None, None) - else: - y = pm.MvNormal(**kwargs, dims=("dsize", "ddata", "dsupport")) - assert pmodel.RV_dims["y"] == ("dsize", "ddata", "dsupport") + y = pm.MvNormal(**kwargs, dims=("dsize", "ddata", "dsupport")) + assert pmodel.RV_dims["y"] == ("dsize", "ddata", "dsupport") assert "dsize" in pmodel.dim_lengths assert y.eval().shape == (2, 3, 4) + def test_simultaneous_dims_and_observed(self): + with pm.Model() as pmodel: + x = pm.ConstantData("x", [1, 2, 3], dims="ddata") + assert "ddata" in pmodel.dim_lengths + + # Note: No checks are performed that observed and dims actually match. + y = pm.Normal("y", observed=[0, 0, 0], dims="ddata") + assert pmodel.RV_dims["y"] == ("ddata",) + assert y.eval().shape == (3,) + def test_define_dims_on_the_fly(self): with pm.Model() as pmodel: agedata = aesara.shared(np.array([10, 20, 30])) @@ -331,11 +320,26 @@ def test_define_dims_on_the_fly(self): # The change should propagate all the way through assert effect.eval().shape == (4,) + def test_define_dims_on_the_fly_from_observed(self): + with pm.Model() as pmodel: + data = aesara.shared(np.zeros((4, 5))) + x = pm.Normal("x", observed=data, dims=("patient", "trials")) + assert pmodel.dim_lengths["patient"].eval() == 4 + assert pmodel.dim_lengths["trials"].eval() == 5 + + # Use dim to create a new RV + x_noisy = pm.Normal("x_noisy", 0, dims=("patient", "trials")) + assert x_noisy.eval().shape == (4, 5) + + # Change data patient dims + data.set_value(np.zeros((10, 6))) + assert x_noisy.eval().shape == (10, 6) + def test_can_resize_data_defined_size(self): with pm.Model() as pmodel: x = pm.MutableData("x", [[1, 2, 3, 4]], dims=("first", "second")) y = pm.Normal("y", mu=0, dims=("first", "second")) - z = pm.Normal("z", mu=y, observed=np.ones((1, 4))) + z = pm.Normal("z", mu=y, observed=np.ones((1, 4)), size=y.shape) assert x.eval().shape == (1, 4) assert y.eval().shape == (1, 4) assert z.eval().shape == (1, 4) @@ -382,7 +386,6 @@ def test_dist_api_works(self): pm.Normal.dist(mu=mu, dims=("town",)) assert pm.Normal.dist(mu=mu, shape=(3,)).eval().shape == (3,) assert pm.Normal.dist(mu=mu, shape=(5, 3)).eval().shape == (5, 3) - assert pm.Normal.dist(mu=mu, shape=(7, ...)).eval().shape == (7, 3) assert pm.Normal.dist(mu=mu, size=(3,)).eval().shape == (3,) assert pm.Normal.dist(mu=mu, size=(4, 3)).eval().shape == (4, 3) @@ -408,10 +411,6 @@ def test_mvnormal_shape_size_difference(self): assert rv.ndim == 3 assert tuple(rv.shape.eval()) == (5, 4, 3) - rv = pm.MvNormal.dist(mu=np.ones((4, 3, 2)), cov=np.eye(2), shape=(6, 5, ...)) - assert rv.ndim == 5 - assert tuple(rv.shape.eval()) == (6, 5, 4, 3, 2) - rv = pm.MvNormal.dist(mu=[1, 2, 3], cov=np.eye(3), size=(5, 4)) assert tuple(rv.shape.eval()) == (5, 4, 3) @@ -422,22 +421,16 @@ def test_convert_dims(self): assert convert_dims(dims="town") == ("town",) with pytest.raises(ValueError, match="must be a tuple, str or list"): convert_dims(3) - with pytest.raises(ValueError, match="may only appear in the last position"): - convert_dims(dims=(..., "town")) def test_convert_shape(self): assert convert_shape(5) == (5,) with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): convert_shape(shape="notashape") - with pytest.raises(ValueError, match="may only appear in the last position"): - convert_shape(shape=(3, ..., 2)) def test_convert_size(self): assert convert_size(7) == (7,) with pytest.raises(ValueError, match="tuple, TensorVariable, int or list"): convert_size(size="notasize") - with pytest.raises(ValueError, match="cannot contain"): - convert_size(size=(3, ...)) def test_lazy_flavors(self): assert pm.Uniform.dist(2, [4, 5], size=[3, 2]).eval().shape == (3, 2) diff --git a/pymc/tests/test_shared.py b/pymc/tests/test_shared.py index 5d92239060..ce91c93c21 100644 --- a/pymc/tests/test_shared.py +++ b/pymc/tests/test_shared.py @@ -41,7 +41,7 @@ def test_sample(self): with pm.Model() as model: 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)