From 6429b0da722d2d4da8fa8a02ce4c64f493426428 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 4 Dec 2023 12:53:35 +0100 Subject: [PATCH 1/3] Allow defining mode from compile_* functions --- pymc/model/core.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 6cfece01cf..cc25366a75 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -612,6 +612,7 @@ def compile_logp( vars: Optional[Union[Variable, Sequence[Variable]]] = None, jacobian: bool = True, sum: bool = True, + **compile_kwargs, ) -> PointFunc: """Compiled log probability density function. @@ -626,12 +627,13 @@ def compile_logp( Whether to sum all logp terms or return elemwise logp for each variable. Defaults to True. """ - return self.compile_fn(self.logp(vars=vars, jacobian=jacobian, sum=sum)) + return self.compile_fn(self.logp(vars=vars, jacobian=jacobian, sum=sum), **compile_kwargs) def compile_dlogp( self, vars: Optional[Union[Variable, Sequence[Variable]]] = None, jacobian: bool = True, + **compile_kwargs, ) -> PointFunc: """Compiled log probability density gradient function. @@ -643,12 +645,13 @@ def compile_dlogp( jacobian : bool Whether to include jacobian terms in logprob graph. Defaults to True. """ - return self.compile_fn(self.dlogp(vars=vars, jacobian=jacobian)) + return self.compile_fn(self.dlogp(vars=vars, jacobian=jacobian), **compile_kwargs) def compile_d2logp( self, vars: Optional[Union[Variable, Sequence[Variable]]] = None, jacobian: bool = True, + **compile_kwargs, ) -> PointFunc: """Compiled log probability density hessian function. @@ -660,7 +663,7 @@ def compile_d2logp( jacobian : bool Whether to include jacobian terms in logprob graph. Defaults to True. """ - return self.compile_fn(self.d2logp(vars=vars, jacobian=jacobian)) + return self.compile_fn(self.d2logp(vars=vars, jacobian=jacobian), **compile_kwargs) def logp( self, From bad219a7821f12e066fbef594e4c2d6d33aea2ed Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 4 Dec 2023 11:18:59 +0100 Subject: [PATCH 2/3] Make coords and data always mutable --- pymc/data.py | 71 +--- pymc/gp/hsgp_approx.py | 4 +- pymc/model/core.py | 36 +- pymc/model/fgraph.py | 7 +- pymc/model_graph.py | 17 +- pymc/sampling/forward.py | 6 +- tests/backends/test_arviz.py | 26 +- tests/backends/test_mcbackend.py | 8 +- tests/distributions/test_shape_utils.py | 8 +- tests/distributions/test_timeseries.py | 4 +- tests/model/test_core.py | 376 +++++++++------------ tests/model/test_fgraph.py | 14 +- tests/model/transform/test_basic.py | 6 +- tests/model/transform/test_conditioning.py | 6 +- tests/sampling/test_forward.py | 56 ++- tests/sampling/test_jax.py | 6 +- tests/sampling/test_mcmc_external.py | 6 +- tests/test_data.py | 73 ++-- tests/test_model_graph.py | 12 +- tests/test_printing.py | 12 +- tests/variational/test_minibatch_rv.py | 4 +- 21 files changed, 318 insertions(+), 440 deletions(-) diff --git a/pymc/data.py b/pymc/data.py index ae2a960421..576cad6b11 100644 --- a/pymc/data.py +++ b/pymc/data.py @@ -262,21 +262,18 @@ def ConstantData( *, dims: Optional[Sequence[str]] = None, coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None, - export_index_as_coords=False, infer_dims_and_coords=False, **kwargs, ) -> TensorConstant: - """Alias for ``pm.Data(..., mutable=False)``. + """Alias for ``pm.Data``. Registers the ``value`` as a :class:`~pytensor.tensor.TensorConstant` with the model. For more information, please reference :class:`pymc.Data`. """ - if export_index_as_coords: - infer_dims_and_coords = export_index_as_coords - warnings.warn( - "Deprecation warning: 'export_index_as_coords; is deprecated and will be removed in future versions. Please use 'infer_dims_and_coords' instead.", - DeprecationWarning, - ) + warnings.warn( + "ConstantData is deprecated. All Data variables are now mutable. Use Data instead.", + FutureWarning, + ) var = Data( name, @@ -284,7 +281,6 @@ def ConstantData( dims=dims, coords=coords, infer_dims_and_coords=infer_dims_and_coords, - mutable=False, **kwargs, ) return cast(TensorConstant, var) @@ -296,21 +292,18 @@ def MutableData( *, dims: Optional[Sequence[str]] = None, coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None, - export_index_as_coords=False, infer_dims_and_coords=False, **kwargs, ) -> SharedVariable: - """Alias for ``pm.Data(..., mutable=True)``. + """Alias for ``pm.Data``. Registers the ``value`` as a :class:`~pytensor.compile.sharedvalue.SharedVariable` with the model. For more information, please reference :class:`pymc.Data`. """ - if export_index_as_coords: - infer_dims_and_coords = export_index_as_coords - warnings.warn( - "Deprecation warning: 'export_index_as_coords; is deprecated and will be removed in future versions. Please use 'infer_dims_and_coords' instead.", - DeprecationWarning, - ) + warnings.warn( + "MutableData is deprecated. All Data variables are now mutable. Use Data instead.", + FutureWarning, + ) var = Data( name, @@ -318,7 +311,6 @@ def MutableData( dims=dims, coords=coords, infer_dims_and_coords=infer_dims_and_coords, - mutable=True, **kwargs, ) return cast(SharedVariable, var) @@ -330,7 +322,6 @@ def Data( *, dims: Optional[Sequence[str]] = None, coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None, - export_index_as_coords=False, infer_dims_and_coords=False, mutable: Optional[bool] = None, **kwargs, @@ -373,15 +364,6 @@ def Data( infer_dims_and_coords : bool, default=False If True, the ``Data`` container will try to infer what the coordinates and dimension names should be if there is an index in ``value``. - mutable : bool, optional - Switches between creating a :class:`~pytensor.compile.sharedvalue.SharedVariable` - (``mutable=True``) vs. creating a :class:`~pytensor.tensor.TensorConstant` - (``mutable=False``). - Consider using :class:`pymc.ConstantData` or :class:`pymc.MutableData` as less - verbose alternatives to ``pm.Data(..., mutable=...)``. - If this parameter is not specified, the value it takes will depend on the - version of the package. Since ``v4.1.0`` the default value is - ``mutable=False``, with previous versions having ``mutable=True``. **kwargs : dict, optional Extra arguments passed to :func:`pytensor.shared`. @@ -394,7 +376,7 @@ def Data( >>> observed_data = [mu + np.random.randn(20) for mu in true_mu] >>> with pm.Model() as model: - ... data = pm.MutableData('data', observed_data[0]) + ... data = pm.Data('data', observed_data[0]) ... mu = pm.Normal('mu', 0, 10) ... pm.Normal('y', mu=mu, sigma=1, observed=data) @@ -430,19 +412,12 @@ def Data( "Pass them directly to `observed` if you want to trigger auto-imputation" ) - if mutable is None: + if mutable is not None: warnings.warn( - "The `mutable` kwarg was not specified. Before v4.1.0 it defaulted to `pm.Data(mutable=True)`," - " which is equivalent to using `pm.MutableData()`." - " In v4.1.0 the default changed to `pm.Data(mutable=False)`, equivalent to `pm.ConstantData`." - " Use `pm.ConstantData`/`pm.MutableData` or pass `pm.Data(..., mutable=False/True)` to avoid this warning.", - UserWarning, + "Data is now always mutable. Specifying the `mutable` kwarg will raise an error in a future release", + FutureWarning, ) - mutable = False - if mutable: - x = pytensor.shared(arr, name, **kwargs) - else: - x = pt.as_tensor_variable(arr, name, **kwargs) + x = pytensor.shared(arr, name, **kwargs) if isinstance(dims, str): dims = (dims,) @@ -453,24 +428,11 @@ def Data( expected=x.ndim, ) - # Optionally infer coords and dims from the input value. - if export_index_as_coords: - infer_dims_and_coords = export_index_as_coords - warnings.warn( - "Deprecation warning: 'export_index_as_coords; is deprecated and will be removed in future versions. Please use 'infer_dims_and_coords' instead.", - DeprecationWarning, - ) - if infer_dims_and_coords: coords, dims = determine_coords(model, value, dims) if dims: - if not mutable: - # Use the dimension lengths from the before it was tensorified. - # These can still be tensors, but in many cases they are numeric. - xshape = np.shape(arr) - else: - xshape = x.shape + xshape = x.shape # Register new dimension lengths for d, dname in enumerate(dims): if dname not in model.dim_lengths: @@ -479,7 +441,6 @@ def Data( # Note: Coordinate values can't be taken from # the value, because it could be N-dimensional. values=coords.get(dname, None), - mutable=mutable, length=xshape[d], ) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index d2675171ee..2778fd370e 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -252,7 +252,7 @@ def prior_linearized(self, Xs: TensorLike): eigenfunctions `phi`, and the square root of the power spectral density. Correct results when using `prior_linearized` in tandem with `pm.set_data` and - `pm.MutableData` require two conditions. First, one must specify `L` instead of `c` when + `pm.Data` require two conditions. First, one must specify `L` instead of `c` when the GP is constructed. If not, a RuntimeError is raised. Second, the `Xs` needs to be zero-centered, so its mean must be subtracted. An example is given below. @@ -290,7 +290,7 @@ def prior_linearized(self, Xs: TensorLike): # First calculate the mean, then make X a shared variable, then subtract the mean. # When X is mutated later, the correct mean will be subtracted. X_mean = np.mean(X, axis=0) - X = pm.MutableData("X", X) + X = pm.Data("X", X) Xs = X - X_mean # Pass the zero-subtracted Xs in to the GP diff --git a/pymc/model/core.py b/pymc/model/core.py index cc25366a75..f117bb1b92 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -515,6 +515,12 @@ def __init__( self.name = self._validate_name(name) self.check_bounds = check_bounds + if coords_mutable is not None: + warnings.warn( + "All coords are now mutable by default. coords_mutable will be removed in a future release.", + FutureWarning, + ) + if self.parent is not None: self.named_vars = treedict(parent=self.parent.named_vars) self.named_vars_to_dims = treedict(parent=self.parent.named_vars_to_dims) @@ -951,7 +957,7 @@ def add_coord( self, name: str, values: Optional[Sequence] = None, - mutable: bool = False, + mutable: Optional[bool] = None, *, length: Optional[Union[int, Variable]] = None, ): @@ -972,6 +978,12 @@ def add_coord( A scalar of the dimensions length. Defaults to ``pytensor.tensor.constant(len(values))``. """ + if mutable is not None: + warnings.warn( + "Coords are now always mutable. Specifying `mutable` will raise an error in a future release", + FutureWarning, + ) + if name in {"draw", "chain", "__sample__"}: raise ValueError( "Dimensions can not be named `draw`, `chain` or `__sample__`, " @@ -995,10 +1007,7 @@ def add_coord( if length is None: length = len(values) if not isinstance(length, Variable): - if mutable: - length = pytensor.shared(length, name=name) - else: - length = pytensor.tensor.constant(length) + length = pytensor.shared(length, name=name) assert length.type.ndim == 0 self._dim_lengths[name] = length self._coords[name] = values @@ -1029,8 +1038,6 @@ def set_dim(self, name: str, new_length: int, coord_values: Optional[Sequence] = coord_values : array_like, optional Optional sequence of coordinate values. """ - if not isinstance(self.dim_lengths[name], SharedVariable): - raise ValueError(f"The dimension '{name}' is immutable.") if coord_values is None and self.coords.get(name, None) is not None: raise ValueError( f"'{name}' has coord values. Pass `set_dim(..., coord_values=...)` to update them." @@ -1079,7 +1086,7 @@ def set_data( ): """Changes the values of a data variable in the model. - In contrast to pm.MutableData().set_value, this method can also + In contrast to pm.Data().set_value, this method can also update the corresponding coordinates. Parameters @@ -1097,7 +1104,7 @@ def set_data( if not isinstance(shared_object, SharedVariable): raise TypeError( f"The variable `{name}` must be a `SharedVariable`" - " (created through `pm.MutableData()` or `pm.Data(mutable=True)`) to allow updating. " + " (created through `pm.Data()` or `pm.Data(mutable=True)`) to allow updating. " f"The current type is: {type(shared_object)}" ) @@ -1114,7 +1121,8 @@ def set_data( for d, dname in enumerate(dims): length_tensor = self.dim_lengths[dname] - old_length = length_tensor.eval() + with pytensor.config.change_flags(cxx=""): + old_length = length_tensor.eval() new_length = values.shape[d] original_coords = self.coords.get(dname, None) new_coords = coords.get(dname, None) @@ -1122,7 +1130,7 @@ def set_data( length_changed = new_length != old_length # Reject resizing if we already know that it would create shape problems. - # NOTE: If there are multiple pm.MutableData containers sharing this dim, but the user only + # NOTE: If there are multiple pm.Data containers sharing this dim, but the user only # changes the values for one of them, they will run into shape problems nonetheless. if length_changed: if original_coords is not None: @@ -1984,8 +1992,8 @@ def set_data(new_data, model=None, *, coords=None): import pymc as pm with pm.Model() as model: - x = pm.MutableData('x', [1., 2., 3.]) - y = pm.MutableData('y', [1., 2., 3.]) + x = pm.Data('x', [1., 2., 3.]) + y = pm.Data('y', [1., 2., 3.]) beta = pm.Normal('beta', 0, 1) obs = pm.Normal('obs', x * beta, 1, observed=y, shape=x.shape) idata = pm.sample() @@ -2014,7 +2022,7 @@ def set_data(new_data, model=None, *, coords=None): data = rng.normal(loc=1.0, scale=2.0, size=100) with pm.Model() as model: - y = pm.MutableData('y', data) + y = pm.Data('y', data) theta = pm.Normal('theta', mu=0.0, sigma=10.0) obs = pm.Normal('obs', theta, 2.0, observed=y, shape=y.shape) idata = pm.sample() diff --git a/pymc/model/fgraph.py b/pymc/model/fgraph.py index 48903c9b72..e2660f210d 100644 --- a/pymc/model/fgraph.py +++ b/pymc/model/fgraph.py @@ -182,8 +182,7 @@ def fgraph_from_model( for named_val in named_value_vars: idx = value_vars.index(named_val) value_vars[idx] = named_val - # Other variables that are in named_vars but are not any of the categories above - # E.g., MutableData, ConstantData, _dim_lengths + # Other variables that are in named_vars but are not any of the categories above (e.g., Data) # We use the same trick as deterministics! accounted_for = set(free_rvs + observed_rvs + potentials + old_deterministics + old_value_vars) other_named_vars = [ @@ -200,8 +199,8 @@ def fgraph_from_model( # Replace the following shared variables in the model: # 1. RNGs - # 2. MutableData (could increase memory usage significantly) - # 3. Mutable coords dim lengths + # 2. Data (could increase memory usage significantly) + # 3. Symbolic coords dim lengths shared_vars_to_copy = find_rng_nodes(model_vars) shared_vars_to_copy += [v for v in model.dim_lengths.values() if isinstance(v, SharedVariable)] shared_vars_to_copy += [v for v in model.named_vars.values() if isinstance(v, SharedVariable)] diff --git a/pymc/model_graph.py b/pymc/model_graph.py index 7f356eb479..d9189d552f 100644 --- a/pymc/model_graph.py +++ b/pymc/model_graph.py @@ -19,14 +19,13 @@ from typing import Optional from pytensor import function -from pytensor.compile.sharedvalue import SharedVariable from pytensor.graph import Apply from pytensor.graph.basic import ancestors, walk from pytensor.scalar.basic import Cast from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.random.op import RandomVariable from pytensor.tensor.shape import Shape -from pytensor.tensor.variable import TensorConstant, TensorVariable +from pytensor.tensor.variable import TensorVariable import pymc as pm @@ -162,14 +161,6 @@ def _make_node(self, var_name, graph, *, nx=False, cluster=False, formatting: st shape = "octagon" style = "filled" label = f"{var_name}\n~\nPotential" - elif isinstance(v, TensorConstant): - shape = "box" - style = "rounded, filled" - label = f"{var_name}\n~\nConstantData" - elif isinstance(v, SharedVariable): - shape = "box" - style = "rounded, filled" - label = f"{var_name}\n~\nMutableData" elif v in self.model.basic_RVs: shape = "ellipse" if v in self.model.observed_RVs: @@ -180,10 +171,14 @@ def _make_node(self, var_name, graph, *, nx=False, cluster=False, formatting: st if symbol.endswith("RV"): symbol = symbol[:-2] label = f"{var_name}\n~\n{symbol}" - else: + elif v in self.model.deterministics: shape = "box" style = None label = f"{var_name}\n~\nDeterministic" + else: + shape = "box" + style = "rounded, filled" + label = f"{var_name}\n~\nCData" kwargs = { "shape": shape, diff --git a/pymc/sampling/forward.py b/pymc/sampling/forward.py index 3a1fc5a785..bf78317911 100644 --- a/pymc/sampling/forward.py +++ b/pymc/sampling/forward.py @@ -115,7 +115,7 @@ def compile_forward_sampling_function( Concretely, this function can be used to compile a function to sample from the posterior predictive distribution of a model that has variables that are conditioned - on ``MutableData`` instances. The variables that depend on the mutable data that have changed + on ``Data`` instances. The variables that depend on the mutable data that have changed will be considered volatile, and as such, they wont be included as inputs into the compiled function. This means that if they have values stored in the posterior, these values will be ignored and new values will be computed (in the case of deterministics and potentials) or @@ -147,8 +147,8 @@ def compile_forward_sampling_function( in the compiled function. The types of the key and value should match or an error will be raised during compilation. constant_data : Optional[Dict[str, numpy.ndarray]] - A dictionary that maps the names of ``MutableData`` or ``ConstantData`` instances to their - corresponding values at inference time. If a model was created with ``MutableData``, these + A dictionary that maps the names of ``Data`` instances to their + corresponding values at inference time. If a model was created with ``Data``, these are stored as ``SharedVariable`` with the name of the data variable and a value equal to the initial data. At inference time, this information is stored in ``InferenceData`` objects under the ``constant_data`` group, which allows us to check whether a diff --git a/tests/backends/test_arviz.py b/tests/backends/test_arviz.py index cf9bc3fc00..adce0c6f5b 100644 --- a/tests/backends/test_arviz.py +++ b/tests/backends/test_arviz.py @@ -268,7 +268,7 @@ def test_autodetect_coords_from_model(self, use_context): ) data_dims = ("date", "city") - data = pm.ConstantData("data", df_data, dims=data_dims) + data = pm.Data("data", df_data, dims=data_dims) _ = pm.Normal( "likelihood", mu=city_temperature, sigma=0.5, observed=data, dims=data_dims ) @@ -307,7 +307,7 @@ def test_overwrite_model_coords_dims(self): x_data = np.arange(4).reshape((2, 2)) y = x_data + np.random.normal(size=(2, 2)) with pm.Model(coords=coords): - x = pm.ConstantData("x", x_data, dims=("dim1", "dim2")) + x = pm.Data("x", x_data, dims=("dim1", "dim2")) beta = pm.Normal("beta", 0, 1, dims="dim1") _ = pm.Normal("obs", x * beta, 1, observed=y, dims=("dim1", "dim2")) trace = pm.sample(100, tune=100, return_inferencedata=False) @@ -438,9 +438,9 @@ def test_potential(self): def test_constant_data(self, use_context): """Test constant_data group behaviour.""" with pm.Model() as model: - x = pm.ConstantData("x", [1.0, 2.0, 3.0]) - y = pm.MutableData("y", [1.0, 2.0, 3.0]) - beta_sigma = pm.MutableData("beta_sigma", 1) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) + beta_sigma = pm.Data("beta_sigma", 1) beta = pm.Normal("beta", 0, beta_sigma) obs = pm.Normal("obs", x * beta, 1, observed=y) trace = pm.sample(100, chains=2, tune=100, return_inferencedata=False) @@ -462,8 +462,8 @@ def test_constant_data(self, use_context): def test_predictions_constant_data(self): with pm.Model(): - x = pm.ConstantData("x", [1.0, 2.0, 3.0]) - y = pm.MutableData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 1) obs = pm.Normal("obs", x * beta, 1, observed=y) trace = pm.sample(100, tune=100, return_inferencedata=False) @@ -474,8 +474,8 @@ def test_predictions_constant_data(self): assert not fails with pm.Model(): - x = pm.MutableData("x", [1.0, 2.0]) - y = pm.ConstantData("y", [1.0, 2.0]) + x = pm.Data("x", [1.0, 2.0]) + y = pm.Data("y", [1.0, 2.0]) beta = pm.Normal("beta", 0, 1) obs = pm.Normal("obs", x * beta, 1, observed=y) predictive_trace = pm.sample_posterior_predictive( @@ -502,8 +502,8 @@ def test_predictions_constant_data(self): def test_no_trace(self): with pm.Model() as model: - x = pm.ConstantData("x", [1.0, 2.0, 3.0]) - y = pm.MutableData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 1) obs = pm.Normal("obs", x * beta, 1, observed=y) idata = pm.sample(100, tune=100) @@ -536,8 +536,8 @@ def test_no_trace(self): def test_priors_separation(self, use_context): """Test model is enough to get prior, prior predictive, constant_data and observed_data.""" with pm.Model() as model: - x = pm.MutableData("x", [1.0, 2.0, 3.0]) - y = pm.ConstantData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 1) obs = pm.Normal("obs", x * beta, 1, observed=y) prior = pm.sample_prior_predictive(return_inferencedata=False) diff --git a/tests/backends/test_mcbackend.py b/tests/backends/test_mcbackend.py index cf80a446d1..6d0704d1a5 100644 --- a/tests/backends/test_mcbackend.py +++ b/tests/backends/test_mcbackend.py @@ -46,12 +46,12 @@ def simple_model(): "condition": ["A", "B", "C"], } ) as pmodel: - x = pm.ConstantData("seconds", seconds, dims="time") + x = pm.Data("seconds", seconds, dims="time") a = pm.Normal("scalar") b = pm.Uniform("vector", dims="condition") pm.Deterministic("matrix", a + b[:, None] * x[None, :], dims=("condition", "time")) pm.Bernoulli("integer", p=0.5) - obs = pm.MutableData("obs", observations, dims=("condition", "time")) + obs = pm.Data("obs", observations, dims=("condition", "time")) pm.Normal("L", pmodel["matrix"], observed=obs, dims=("condition", "time")) return pmodel @@ -65,7 +65,7 @@ def test_find_data(simple_model): assert isinstance(secs, mcb.DataVariable) assert secs.dims == ["time"] assert not secs.is_observed - np.testing.assert_array_equal(ndarray_to_numpy(secs.value), simple_model["seconds"].data) + np.testing.assert_array_equal(ndarray_to_numpy(secs.value), simple_model["seconds"].get_value()) obs = dvardict["obs"] assert isinstance(obs, mcb.DataVariable) @@ -77,7 +77,7 @@ def test_find_data(simple_model): def test_find_data_skips_deterministics(): data = np.array([0, 1], dtype="float32") with pm.Model() as pmodel: - a = pm.ConstantData("a", data, dims="item") + a = pm.Data("a", data, dims="item") b = pm.Normal("b") pm.Deterministic("c", a + b, dims="item") assert "c" in pmodel.named_vars diff --git a/tests/distributions/test_shape_utils.py b/tests/distributions/test_shape_utils.py index 9b16031c53..e6be429c99 100644 --- a/tests/distributions/test_shape_utils.py +++ b/tests/distributions/test_shape_utils.py @@ -188,7 +188,7 @@ def test_broadcast_by_observed(self): def test_simultaneous_shape_and_dims(self): with pm.Model() as pmodel: - x = pm.ConstantData("x", [1, 2, 3], dims="ddata") + x = pm.Data("x", [1, 2, 3], dims="ddata") # The shape and dims tuples correspond to each other. # Note: No checks are performed that implied shape (x), shape and dims actually match. @@ -200,7 +200,7 @@ def test_simultaneous_shape_and_dims(self): def test_simultaneous_size_and_dims(self): with pm.Model() as pmodel: - x = pm.ConstantData("x", [1, 2, 3], dims="ddata") + x = pm.Data("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. @@ -213,7 +213,7 @@ def test_simultaneous_size_and_dims(self): def test_simultaneous_dims_and_observed(self): with pm.Model() as pmodel: - x = pm.ConstantData("x", [1, 2, 3], dims="ddata") + x = pm.Data("x", [1, 2, 3], dims="ddata") assert "ddata" in pmodel.dim_lengths # Note: No checks are performed that observed and dims actually match. @@ -234,7 +234,7 @@ def test_define_dims_on_the_fly_raises(self): def test_can_resize_data_defined_size(self): with pm.Model() as pmodel: - x = pm.MutableData("x", [[1, 2, 3, 4]], dims=("first", "second")) + x = pm.Data("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)), size=y.shape) assert x.eval().shape == (1, 4) diff --git a/tests/distributions/test_timeseries.py b/tests/distributions/test_timeseries.py index 4e1bfd723b..683838435a 100644 --- a/tests/distributions/test_timeseries.py +++ b/tests/distributions/test_timeseries.py @@ -20,7 +20,7 @@ import pymc as pm -from pymc import MutableData +from pymc import Data from pymc.distributions.continuous import Exponential, Flat, HalfNormal, Normal, Uniform from pymc.distributions.distribution import DiracDelta from pymc.distributions.multivariate import ( @@ -406,7 +406,7 @@ def test_gaussian_inference(self): _mu = Uniform("mu", -10, 10) _sigma = Uniform("sigma", 0, 10) - obs_data = MutableData("obs_data", obs) + obs_data = Data("obs_data", obs) grw = GaussianRandomWalk( "grw", _mu, _sigma, steps=steps, observed=obs_data, init_dist=Normal.dist(0, 100) ) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index b1c1b42f57..d62890ab1d 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -33,9 +33,7 @@ from pytensor.graph import graph_inputs from pytensor.raise_op import Assert -from pytensor.tensor import TensorVariable from pytensor.tensor.random.op import RandomVariable -from pytensor.tensor.sharedvar import TensorSharedVariable from pytensor.tensor.variable import TensorConstant import pymc as pm @@ -641,7 +639,7 @@ def test_eval_rv_shapes(self): "city": ["Sydney", "Las Vegas", "Düsseldorf"], } ) as pmodel: - pm.MutableData("budget", [1, 2, 3, 4], dims="year") + pm.Data("budget", [1, 2, 3, 4], dims="year") pm.Normal("untransformed", size=(1, 2)) pm.Uniform("transformed", size=(7,)) obs = pm.Uniform("observed", size=(3,), observed=[0.1, 0.2, 0.3]) @@ -746,232 +744,174 @@ def test_nested_model_coords(): assert set(m2.named_vars_to_dims) < set(m1.named_vars_to_dims) -def test_shapeerror_from_set_data_dimensionality(): - with pm.Model() as pmodel: - pm.MutableData("m", np.ones((3,)), dims="one") - with pytest.raises(ValueError, match="must have 1 dimensions"): - pmodel.set_data("m", np.ones((3, 4))) - - -def test_shapeerror_from_resize_immutable_dim_from_RV(): - """ - Trying to resize an immutable dimension should raise a ShapeError. - Even if the variable being updated is a SharedVariable and has other - dimensions that are mutable. - """ - with pm.Model(coords={"fixed": range(3)}) as pmodel: - pm.Normal("a", mu=[1, 2, 3], dims="fixed") - assert isinstance(pmodel.dim_lengths["fixed"], TensorVariable) - - pm.MutableData("m", [[1, 2, 3]], dims=("one", "fixed")) - - # This is fine because the "fixed" dim is not resized - pmodel.set_data("m", [[1, 2, 3], [3, 4, 5]]) - - msg = "The 'm' variable already had 3 coord values defined for its fixed dimension" - with pytest.raises(ValueError, match=msg): - # Can't work because the "fixed" dimension is linked to a - # TensorVariable with constant shape. - # Note that the new data tries to change both dimensions - pmodel.set_data("m", [[1, 2], [3, 4]]) - - -def test_shapeerror_from_resize_immutable_dim_from_coords(): - with pm.Model(coords={"immutable": [1, 2]}) as pmodel: - assert isinstance(pmodel.dim_lengths["immutable"], TensorConstant) - pm.MutableData("m", [1, 2], dims="immutable") - # Data can be changed - pmodel.set_data("m", [3, 4]) - - with pytest.raises(ShapeError, match="`TensorConstant` stores its length"): - # But the length is linked to a TensorConstant - pmodel.set_data("m", [1, 2, 3], coords=dict(immutable=[1, 2, 3])) - - -def test_valueerror_from_resize_without_coords_update(): - """ - Resizing a mutable dimension that had coords, - without passing new coords raises a ValueError. - """ - with pm.Model() as pmodel: - pmodel.add_coord("shared", [1, 2, 3], mutable=True) - pm.MutableData("m", [1, 2, 3], dims=("shared")) - with pytest.raises(ValueError, match="'m' variable already had 3"): - # tries to resize m but without passing coords so raise ValueError - pm.set_data({"m": [1, 2, 3, 4]}) - - -def test_coords_and_constantdata_create_immutable_dims(): - """ - When created from `pm.Model(coords=...)` or `pm.ConstantData` - a dimension should be resizable. - """ - with pm.Model(coords={"group": ["A", "B"]}) as m: - x = pm.ConstantData("x", [0], dims="feature") - y = pm.Normal("y", x, 1, dims=("group", "feature")) - assert isinstance(m._dim_lengths["feature"], TensorConstant) - assert isinstance(m._dim_lengths["group"], TensorConstant) - assert x.eval().shape == (1,) - assert y.eval().shape == (2, 1) - - -def test_add_coord_mutable_kwarg(): - """ - Checks resulting tensor type depending on mutable kwarg in add_coord. - """ - with pm.Model() as m: - m.add_coord("fixed", values=[1], mutable=False) - m.add_coord("mutable1", values=[1, 2], mutable=True) - assert isinstance(m._dim_lengths["fixed"], TensorConstant) - assert isinstance(m._dim_lengths["mutable1"], TensorSharedVariable) - pm.MutableData("mdata", np.ones((1, 2, 3)), dims=("fixed", "mutable1", "mutable2")) - assert isinstance(m._dim_lengths["mutable2"], TensorVariable) - - -def test_set_dim(): - """Test the conscious re-sizing of dims created through add_coord().""" - with pm.Model() as pmodel: - pmodel.add_coord("fdim", mutable=False, length=1) - pmodel.add_coord("mdim", mutable=True, length=2) - a = pm.Normal("a", dims="mdim") - assert a.eval().shape == (2,) - - with pytest.raises(ValueError, match="is immutable"): - pmodel.set_dim("fdim", 3) - - pmodel.set_dim("mdim", 3) - assert a.eval().shape == (3,) - - -def test_set_dim_with_coords(): - """Test the conscious re-sizing of dims created through add_coord() with coord value.""" - with pm.Model() as pmodel: - pmodel.add_coord("mdim", mutable=True, length=2, values=["A", "B"]) - a = pm.Normal("a", dims="mdim") - assert len(pmodel.coords["mdim"]) == 2 - - with pytest.raises(ValueError, match="has coord values"): - pmodel.set_dim("mdim", new_length=3) - - with pytest.raises(ShapeError, match="does not match"): - pmodel.set_dim("mdim", new_length=3, coord_values=["A", "B"]) - - pmodel.set_dim("mdim", 3, ["A", "B", "C"]) - assert a.eval().shape == (3,) - assert pmodel.coords["mdim"] == ("A", "B", "C") - - -def test_add_named_variable_checks_dim_name(): - with pm.Model() as pmodel: - rv = pm.Normal.dist(mu=[1, 2]) - - # Checks that vars are named - with pytest.raises(ValueError, match="is unnamed"): - pmodel.add_named_variable(rv) - rv.name = "nomnom" - - # Coords must be available already - with pytest.raises(ValueError, match="not specified in `coords`"): - pmodel.add_named_variable(rv, dims="nomnom") - pmodel.add_coord("nomnom", [1, 2]) - - # No name collisions - with pytest.raises(ValueError, match="same name as"): - pmodel.add_named_variable(rv, dims="nomnom") - - # This should work (regression test against #6335) - rv2 = rv[:, None] - rv2.name = "yumyum" - pmodel.add_named_variable(rv2, dims=("nomnom", None)) - - -def test_dims_type_check(): - with pm.Model(coords={"a": range(5)}) as m: - with pytest.raises(TypeError, match="Dims must be string"): - x = pm.Normal("x", shape=(10, 5), dims=(None, "a")) - - -def test_none_coords_autonumbering(): - with pm.Model() as m: - m.add_coord(name="a", values=None, length=3) - m.add_coord(name="b", values=range(5)) - x = pm.Normal("x", dims=("a", "b")) - prior = pm.sample_prior_predictive(samples=2).prior - assert prior["x"].shape == (1, 2, 3, 5) - assert list(prior.coords["a"].values) == list(range(3)) - assert list(prior.coords["b"].values) == list(range(5)) - - -def test_set_data_indirect_resize(): - with pm.Model() as pmodel: - pmodel.add_coord("mdim", mutable=True, length=2) - pm.MutableData("mdata", [1, 2], dims="mdim") - - # First resize the dimension. - pmodel.dim_lengths["mdim"].set_value(3) - # Then change the data. - with warnings.catch_warnings(): - warnings.simplefilter("error") - pmodel.set_data("mdata", [1, 2, 3]) - - # Now the other way around. - with warnings.catch_warnings(): - warnings.simplefilter("error") - pmodel.set_data("mdata", [1, 2, 3, 4]) - - -def test_set_data_warns_on_resize_of_dims_defined_by_other_mutabledata(): - with pm.Model() as pmodel: - pm.MutableData("m1", [1, 2], dims="mutable") - pm.MutableData("m2", [3, 4], dims="mutable") - - # Resizing the non-defining variable first gives a warning - with pytest.warns(ShapeWarning, match="by another variable"): - pmodel.set_data("m2", [4, 5, 6]) - pmodel.set_data("m1", [1, 2, 3]) - - # Resizing the definint variable first is silent +class TestSetUpdateCoords: + def test_shapeerror_from_set_data_dimensionality(self): + with pm.Model() as pmodel: + pm.Data("m", np.ones((3,)), dims="one") + with pytest.raises(ValueError, match="must have 1 dimensions"): + pmodel.set_data("m", np.ones((3, 4))) + + def test_resize_from_set_data_dim_with_coords(self): + with pm.Model(coords={"dim_with_coords": [1, 2]}) as pmodel: + pm.Data("m", [1, 2], dims=("dim_with_coords",)) + # Does not resize dim + pmodel.set_data("m", [3, 4]) + # Resizes, but also passes new coords + pmodel.set_data("m", [1, 2, 3], coords=dict(dim_with_coords=[1, 2, 3])) + + # Resizes, but does not pass new coords + with pytest.raises(ValueError, match="'m' variable already had 3"): + pm.set_data({"m": [1, 2, 3, 4]}) + + def test_resize_from_set_data_dim_without_coords(self): + with pm.Model() as pmodel: + # TODO: Either support dims without coords everywhere or don't. Why is it okay for Data but not RVs? + pm.Data("m", [1, 2], dims=("dim_without_coords",)) + pmodel.set_data("m", [3, 4]) + pmodel.set_data("m", [1, 2, 3]) + + def test_resize_from_set_dim(self): + """Test the conscious re-sizing of dims created through add_coord() with coord value.""" + with pm.Model(coords={"mdim": ["A", "B"]}) as pmodel: + a = pm.Normal("a", dims="mdim") + assert pmodel.coords["mdim"] == ("A", "B") + + with pytest.raises(ValueError, match="has coord values"): + pmodel.set_dim("mdim", new_length=3) + + with pytest.raises(ShapeError, match="does not match"): + pmodel.set_dim("mdim", new_length=3, coord_values=["A", "B"]) + + pmodel.set_dim("mdim", 3, ["A", "B", "C"]) + assert pmodel.coords["mdim"] == ("A", "B", "C") + with pytensor.config.change_flags(cxx=""): + assert a.eval().shape == (3,) + + def test_resize_from_set_data_and_set_dim(self): + with pm.Model(coords={"group": ["A", "B"]}) as m: + x = pm.Data("x", [0], dims="feature") + y = pm.Normal("y", x, 1, dims=("group", "feature")) + + with pytensor.config.change_flags(cxx=""): + assert x.eval().shape == (1,) + assert y.eval().shape == (2, 1) + + m.set_data("x", [0, 1]) + m.set_dim("group", new_length=3, coord_values=["A", "B", "C"]) + with pytensor.config.change_flags(cxx=""): + assert x.eval().shape == (2,) + assert y.eval().shape == (3, 2) + + def test_add_named_variable_checks_dim_name(self): + with pm.Model() as pmodel: + rv = pm.Normal.dist(mu=[1, 2]) + + # Checks that vars are named + with pytest.raises(ValueError, match="is unnamed"): + pmodel.add_named_variable(rv) + rv.name = "nomnom" + + # Coords must be available already + with pytest.raises(ValueError, match="not specified in `coords`"): + pmodel.add_named_variable(rv, dims="nomnom") + pmodel.add_coord("nomnom", [1, 2]) + + # No name collisions + with pytest.raises(ValueError, match="same name as"): + pmodel.add_named_variable(rv, dims="nomnom") + + # This should work (regression test against #6335) + rv2 = rv[:, None] + rv2.name = "yumyum" + pmodel.add_named_variable(rv2, dims=("nomnom", None)) + + def test_dims_type_check(self): + with pm.Model(coords={"a": range(5)}) as m: + with pytest.raises(TypeError, match="Dims must be string"): + x = pm.Normal("x", shape=(10, 5), dims=(None, "a")) + + def test_none_coords_autonumbering(self): + # TODO: Either allow dims without coords everywhere or nowhere + with pm.Model() as m: + m.add_coord(name="a", values=None, length=3) + m.add_coord(name="b", values=range(5)) + x = pm.Normal("x", dims=("a", "b")) + prior = pm.sample_prior_predictive(samples=2).prior + assert prior["x"].shape == (1, 2, 3, 5) + assert list(prior.coords["a"].values) == list(range(3)) + assert list(prior.coords["b"].values) == list(range(5)) + + def test_set_data_indirect_resize_without_coords(self): + with pm.Model() as pmodel: + pmodel.add_coord("mdim", length=2) + pm.Data("mdata", [1, 2], dims="mdim") + + assert pmodel.dim_lengths["mdim"].get_value() == 2 + assert pmodel.coords["mdim"] is None + + # First resize the dimension. + pmodel.dim_lengths["mdim"].set_value(3) + # Then change the data. with warnings.catch_warnings(): warnings.simplefilter("error") - pmodel.set_data("m1", [1, 2]) - pmodel.set_data("m2", [3, 4]) + pmodel.set_data("mdata", [1, 2, 3]) + # Now the other way around. + with warnings.catch_warnings(): + warnings.simplefilter("error") + pmodel.set_data("mdata", [1, 2, 3, 4]) -def test_set_data_indirect_resize_with_coords(): - with pm.Model() as pmodel: - pmodel.add_coord("mdim", ["A", "B"], mutable=True, length=2) - pm.MutableData("mdata", [1, 2], dims="mdim") + def test_set_data_indirect_resize_with_coords(self): + with pm.Model() as pmodel: + pmodel.add_coord("mdim", ["A", "B"], mutable=True, length=2) + pm.Data("mdata", [1, 2], dims="mdim") - assert pmodel.coords["mdim"] == ("A", "B") + assert pmodel.coords["mdim"] == ("A", "B") - # First resize the dimension. - pmodel.set_dim("mdim", 3, ["A", "B", "C"]) - assert pmodel.coords["mdim"] == ("A", "B", "C") - # Then change the data. - with warnings.catch_warnings(): - warnings.simplefilter("error") - pmodel.set_data("mdata", [1, 2, 3]) + # First resize the dimension. + pmodel.set_dim("mdim", 3, ["A", "B", "C"]) + assert pmodel.coords["mdim"] == ("A", "B", "C") + # Then change the data. + with warnings.catch_warnings(): + warnings.simplefilter("error") + pmodel.set_data("mdata", [1, 2, 3]) - # Now the other way around. - with warnings.catch_warnings(): - warnings.simplefilter("error") - pmodel.set_data("mdata", [1, 2, 3, 4], coords=dict(mdim=["A", "B", "C", "D"])) - assert pmodel.coords["mdim"] == ("A", "B", "C", "D") + # Now the other way around. + with warnings.catch_warnings(): + warnings.simplefilter("error") + pmodel.set_data("mdata", [1, 2, 3, 4], coords=dict(mdim=["A", "B", "C", "D"])) + assert pmodel.coords["mdim"] == ("A", "B", "C", "D") - # This time with incorrectly sized coord values - with pytest.raises(ShapeError, match="new coordinate values"): - pmodel.set_data("mdata", [1, 2], coords=dict(mdim=[1, 2, 3])) + # This time with incorrectly sized coord values + with pytest.raises(ShapeError, match="new coordinate values"): + pmodel.set_data("mdata", [1, 2], coords=dict(mdim=[1, 2, 3])) + def test_set_data_warns_on_resize_of_dims_defined_by_other_data(self): + with pm.Model() as pmodel: + pm.Data("m1", [1, 2], dims="mutable") + pm.Data("m2", [3, 4], dims="mutable") -def test_set_data_constant_shape_error(): - with pm.Model() as pmodel: - x = pm.Normal("x", size=7) - pmodel.add_coord("weekday", length=x.shape[0]) - pm.MutableData("y", np.arange(7), dims="weekday") + # Resizing the non-defining variable first gives a warning + with pytest.warns(ShapeWarning, match="by another variable"): + pmodel.set_data("m2", [4, 5, 6]) + pmodel.set_data("m1", [1, 2, 3]) - msg = "because the dimension was initialized from 'x'" - with pytest.raises(ShapeError, match=msg): - pmodel.set_data("y", np.arange(10)) + # Resizing the defining variable first is silent + with warnings.catch_warnings(): + warnings.simplefilter("error") + pmodel.set_data("m1", [1, 2]) + pmodel.set_data("m2", [3, 4]) + + def test_set_data_constant_shape_error(self): + # TODO: Why allow such a complex scenario? + with pm.Model() as pmodel: + x = pm.Normal("x", size=7) + pmodel.add_coord("weekday", length=x.shape[0]) + pm.Data("y", np.arange(7), dims="weekday") + + msg = "because the dimension was initialized from 'x'" + with pytest.raises(ShapeError, match=msg): + pmodel.set_data("y", np.arange(10)) @pytest.mark.parametrize("jacobian", [True, False]) diff --git a/tests/model/test_fgraph.py b/tests/model/test_fgraph.py index 9580ccd303..0608157ad8 100644 --- a/tests/model/test_fgraph.py +++ b/tests/model/test_fgraph.py @@ -100,15 +100,15 @@ def same_storage(shared_1, shared_2) -> bool: @pytest.mark.parametrize("inline_views", (False, True)) def test_data(inline_views): - """Test shared RNGs, MutableData, ConstantData and dim lengths are handled correctly. + """Test shared RNGs, Data, and dim lengths are handled correctly. All model-related shared variables should be copied to become independent across models. """ - with pm.Model(coords_mutable={"test_dim": range(3)}) as m_old: - x = pm.MutableData("x", [0.0, 1.0, 2.0], dims=("test_dim",)) - y = pm.MutableData("y", [10.0, 11.0, 12.0], dims=("test_dim",)) + with pm.Model(coords={"test_dim": range(3)}) as m_old: + x = pm.Data("x", [0.0, 1.0, 2.0], dims=("test_dim",)) + y = pm.Data("y", [10.0, 11.0, 12.0], dims=("test_dim",)) sigma = pm.MutableData("sigma", [1.0], shape=(1,)) - b0 = pm.ConstantData("b0", np.zeros((1,))) + b0 = pm.Data("b0", np.zeros((1,)), shape=((1,))) b1 = pm.DiracDelta("b1", 1.0) mu = pm.Deterministic("mu", b0 + b1 * x, dims=("test_dim",)) obs = pm.Normal("obs", mu=mu, sigma=sigma, observed=y, dims=("test_dim",)) @@ -139,8 +139,8 @@ def test_data(inline_views): # The rv-data mapping is preserved assert m_new.rvs_to_values[m_new["obs"]] is m_new["y"] - # ConstantData is still accessible as a model variable - np.testing.assert_array_equal(m_new["b0"], m_old["b0"]) + # Data is still accessible as a model variable + np.testing.assert_array_equal(m_new["b0"].get_value(), m_old["b0"].get_value()) # Shared model variables, dim lengths, and rngs are copied and no longer point to the same memory assert not same_storage(m_new["x"], x) diff --git a/tests/model/transform/test_basic.py b/tests/model/transform/test_basic.py index 1328ea6f1c..b62edaafc6 100644 --- a/tests/model/transform/test_basic.py +++ b/tests/model/transform/test_basic.py @@ -18,13 +18,13 @@ def test_prune_vars_detached_from_observed(): with pm.Model() as m: - obs_data = pm.MutableData("obs_data", 0) - a0 = pm.ConstantData("a0", 0) + obs_data = pm.Data("obs_data", 0) + a0 = pm.Data("a0", 0) a1 = pm.Normal("a1", a0) a2 = pm.Normal("a2", a1) pm.Normal("obs", a2, observed=obs_data) - d0 = pm.ConstantData("d0", 0) + d0 = pm.Data("d0", 0) d1 = pm.Normal("d1", d0) assert set(m.named_vars.keys()) == {"obs_data", "a0", "a1", "a2", "obs", "d0", "d1"} diff --git a/tests/model/transform/test_conditioning.py b/tests/model/transform/test_conditioning.py index 724425b68e..a9a8ab712d 100644 --- a/tests/model/transform/test_conditioning.py +++ b/tests/model/transform/test_conditioning.py @@ -132,7 +132,7 @@ def test_do(): # Test two substitutions with m_old: - switch = pm.MutableData("switch", 1) + switch = pm.Data("switch", 1) m_new = do(m_old, {y: 100 * switch, x: 100 * switch}) assert len(m_new.free_RVs) == 1 @@ -213,8 +213,8 @@ def test_do_dims(): @pytest.mark.parametrize("prune", (False, True)) def test_do_prune(prune): with pm.Model() as m: - x0 = pm.ConstantData("x0", 0) - x1 = pm.ConstantData("x1", 0) + x0 = pm.Data("x0", 0) + x1 = pm.Data("x1", 0) y = pm.Normal("y") y_det = pm.Deterministic("y_det", y + x0) z = pm.Normal("z", y_det) diff --git a/tests/sampling/test_forward.py b/tests/sampling/test_forward.py index c901e867e8..0684d13285 100644 --- a/tests/sampling/test_forward.py +++ b/tests/sampling/test_forward.py @@ -124,8 +124,8 @@ def get_function_inputs(function): def test_linear_model(self): with pm.Model() as model: - x = pm.MutableData("x", np.linspace(0, 1, 10)) - y = pm.MutableData("y", np.ones(10)) + x = pm.Data("x", np.linspace(0, 1, 10)) + y = pm.Data("y", np.ones(10)) alpha = pm.Normal("alpha", 0, 0.1) beta = pm.Normal("beta", 0, 0.1) @@ -142,30 +142,11 @@ def test_linear_model(self): assert {i.name for i in self.get_function_inputs(f)} == {"alpha", "beta", "sigma"} assert {i.name for i in self.get_function_roots(f)} == {"x", "alpha", "beta", "sigma"} - with pm.Model() as model: - x = pm.ConstantData("x", np.linspace(0, 1, 10)) - y = pm.MutableData("y", np.ones(10)) - - alpha = pm.Normal("alpha", 0, 0.1) - 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, shape=x.shape) - - f, volatile_rvs = compile_forward_sampling_function( - [obs], - vars_in_trace=[alpha, beta, sigma, mu], - basic_rvs=model.basic_RVs, - ) - assert volatile_rvs == {obs} - assert {i.name for i in self.get_function_inputs(f)} == {"alpha", "beta", "sigma", "mu"} - assert {i.name for i in self.get_function_roots(f)} == {"mu", "sigma"} - def test_nested_observed_model(self): with pm.Model() as model: - p = pm.ConstantData("p", np.array([0.25, 0.5, 0.25])) - x = pm.MutableData("x", np.zeros(10)) - y = pm.MutableData("y", np.ones(10)) + p = pm.Data("p", np.array([0.25, 0.5, 0.25])) + x = pm.Data("x", np.zeros(10)) + y = pm.Data("y", np.ones(10)) category = pm.Categorical("category", p, observed=x) beta = pm.Normal("beta", 0, 0.1, size=p.shape) @@ -178,6 +159,16 @@ def test_nested_observed_model(self): vars_in_trace=[beta, mu, sigma], basic_rvs=model.basic_RVs, ) + assert volatile_rvs == {category, beta, obs} + assert {i.name for i in self.get_function_inputs(f)} == {"sigma"} + assert {i.name for i in self.get_function_roots(f)} == {"x", "p", "sigma"} + + f, volatile_rvs = compile_forward_sampling_function( + outputs=model.observed_RVs, + vars_in_trace=[beta, mu, sigma], + constant_data={"p": p.get_value()}, + basic_rvs=model.basic_RVs, + ) assert volatile_rvs == {category, obs} assert {i.name for i in self.get_function_inputs(f)} == {"beta", "sigma"} assert {i.name for i in self.get_function_roots(f)} == {"x", "p", "beta", "sigma"} @@ -185,6 +176,7 @@ def test_nested_observed_model(self): f, volatile_rvs = compile_forward_sampling_function( outputs=model.observed_RVs, vars_in_trace=[beta, mu, sigma], + constant_data={"p": p.get_value()}, basic_rvs=model.basic_RVs, givens_dict={category: np.zeros(10, dtype=category.dtype)}, ) @@ -200,7 +192,7 @@ def test_nested_observed_model(self): def test_volatile_parameters(self): with pm.Model() as model: - y = pm.MutableData("y", np.ones(10)) + y = pm.Data("y", np.ones(10)) mu = pm.Normal("mu", 0, 1) nested_mu = pm.Normal("nested_mu", mu, 1, size=10) sigma = pm.HalfNormal("sigma", 1) @@ -352,13 +344,13 @@ def test_mutable_coords_volatile(self): with pm.Model() as model: model.add_coord("name", ["A", "B", "C"], mutable=True) model.add_coord("obs", list(range(10, 20)), mutable=True) - offsets = pm.MutableData("offsets", rng.normal(0, 1, size=(10,))) + offsets = pm.Data("offsets", rng.normal(0, 1, size=(10,))) a = pm.Normal("a", mu=0, sigma=1, dims=["name"]) b = pm.Normal("b", mu=offsets, sigma=1) mu = pm.Deterministic("mu", a + b[..., None], dims=["obs", "name"]) sigma = pm.HalfNormal("sigma", sigma=1, dims=["name"]) - data = pm.MutableData( + data = pm.Data( "y_obs", data, dims=["obs", "name"], @@ -883,13 +875,13 @@ def make_mock_model(): with pm.Model() as model: model.add_coord("name", ["A", "B", "C"], mutable=True) model.add_coord("obs", list(range(10, 20)), mutable=True) - offsets = pm.MutableData("offsets", rng.normal(0, 1, size=(10,))) + offsets = pm.Data("offsets", rng.normal(0, 1, size=(10,))) a = pm.Normal("a", mu=0, sigma=1, dims=["name"]) b = pm.Normal("b", mu=offsets, sigma=1) mu = pm.Deterministic("mu", a + b[..., None], dims=["obs", "name"]) sigma = pm.HalfNormal("sigma", sigma=1, dims=["name"]) - data = pm.MutableData( + data = pm.Data( "y_obs", data, dims=["obs", "name"], @@ -938,7 +930,7 @@ def test_logging_sampled_basic_rvs_posterior_mutable(self, mock_sample_results, pm.sample_posterior_predictive(samples) if kind == "MultiTrace": # MultiTrace will only have the actual MCMC posterior samples but no information on - # the MutableData and mutable coordinate values, so it will always assume they are volatile + # the Data and coordinate values, so it will always assume they are volatile # and resample their descendants assert caplog.record_tuples == [ ("pymc.sampling.forward", logging.INFO, "Sampling: [a, b, sigma, y]") @@ -1042,7 +1034,7 @@ def test_ignores_observed(self, seeded_test): observed = np.random.normal(10, 1, size=200) with pm.Model(): # Use a prior that's way off to show we're ignoring the observed variables - observed_data = pm.MutableData("observed_data", observed) + observed_data = pm.Data("observed_data", observed) mu = pm.Normal("mu", mu=-100, sigma=1) positive_mu = pm.Deterministic("positive_mu", np.abs(mu)) z = -1 - positive_mu @@ -1636,7 +1628,7 @@ def test_get_vars_in_point_list(): with pm.Model() as modelB: a = pm.Normal("a", 0, 1) pm.Normal("c", 0, 1) - pm.ConstantData("d", 0) + pm.Data("d", 0) point_list = [{"a": 0, "b": 0, "d": 0}] vars_in_trace = get_vars_in_point_list(point_list, modelB) diff --git a/tests/sampling/test_jax.py b/tests/sampling/test_jax.py index 77121b37ce..5b53f98642 100644 --- a/tests/sampling/test_jax.py +++ b/tests/sampling/test_jax.py @@ -268,8 +268,7 @@ def model_test_idata_kwargs() -> pm.Model: x = pm.Normal("x", shape=(2,), dims=["x_coord"]) _ = pm.Normal("y", x, observed=[0, 0]) _ = pm.Normal("z", 0, 1, dims="z_coord") - pm.ConstantData("constantdata", [1, 2, 3]) - pm.MutableData("mutabledata", 2) + pm.Data("data", [1, 2, 3]) return m @@ -312,8 +311,7 @@ def test_idata_kwargs( assert idata is not None const_data = idata.get("constant_data") assert const_data is not None - assert "constantdata" in const_data - assert "mutabledata" in const_data + assert "data" in const_data if idata_kwargs.get("log_likelihood", False): assert "log_likelihood" in idata diff --git a/tests/sampling/test_mcmc_external.py b/tests/sampling/test_mcmc_external.py index a6c0e16680..01b36c004b 100644 --- a/tests/sampling/test_mcmc_external.py +++ b/tests/sampling/test_mcmc_external.py @@ -16,7 +16,7 @@ import numpy.testing as npt import pytest -from pymc import ConstantData, Model, Normal, sample +from pymc import Data, Model, Normal, sample @pytest.mark.parametrize("nuts_sampler", ["pymc", "nutpie", "blackjax", "numpyro"]) @@ -26,8 +26,8 @@ def test_external_nuts_sampler(recwarn, nuts_sampler): with Model(): x = Normal("x", 100, 5) - y = ConstantData("y", [1, 2, 3, 4]) - ConstantData("z", [100, 190, 310, 405]) + y = Data("y", [1, 2, 3, 4]) + Data("z", [100, 190, 310, 405]) Normal("L", mu=x, sigma=0.1, observed=y) diff --git a/tests/test_data.py b/tests/test_data.py index 363c76d5a2..a70b9c1714 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -37,7 +37,7 @@ class TestData: def test_deterministic(self): data_values = np.array([0.5, 0.4, 5, 2]) with pm.Model() as model: - X = pm.MutableData("X", data_values) + X = pm.Data("X", data_values) pm.Normal("y", 0, 1, observed=X) model.compile_logp()(model.initial_point()) @@ -48,7 +48,7 @@ def test_sample(self, seeded_test): x_pred = np.linspace(-3, 3, 200, dtype="float32") with pm.Model(): - x_shared = pm.MutableData("x_shared", x) + x_shared = pm.Data("x_shared", x) b = pm.Normal("b", 0.0, 10.0) pm.Normal("obs", b * x_shared, np.sqrt(1e-2), observed=y, shape=x_shared.shape) @@ -76,8 +76,8 @@ def test_sample(self, seeded_test): def test_sample_posterior_predictive_after_set_data(self): with pm.Model() as model: - x = pm.MutableData("x", [1.0, 2.0, 3.0]) - y = pm.ConstantData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 10.0) pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y) trace = pm.sample( @@ -101,7 +101,7 @@ def test_sample_posterior_predictive_after_set_data(self): def test_sample_posterior_predictive_after_set_data_with_coords(self): y = np.array([1.0, 2.0, 3.0]) with pm.Model() as model: - x = pm.MutableData("x", [1.0, 2.0, 3.0], dims="obs_id") + x = pm.Data("x", [1.0, 2.0, 3.0], dims="obs_id") beta = pm.Normal("beta", 0, 10.0) pm.Normal("obs", beta * x, np.sqrt(1e-3), observed=y, dims="obs_id") idata = pm.sample( @@ -125,8 +125,8 @@ def test_sample_posterior_predictive_after_set_data_with_coords(self): def test_sample_after_set_data(self): with pm.Model() as model: - x = pm.MutableData("x", [1.0, 2.0, 3.0]) - y = pm.MutableData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 10.0) pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y) pm.sample( @@ -159,8 +159,8 @@ def test_shared_data_as_index(self): See https://github.com/pymc-devs/pymc/issues/3813 """ with pm.Model() as model: - index = pm.MutableData("index", [2, 0, 1, 0, 2]) - y = pm.MutableData("y", [1.0, 2.0, 3.0, 2.0, 1.0]) + index = pm.Data("index", [2, 0, 1, 0, 2]) + y = pm.Data("y", [1.0, 2.0, 3.0, 2.0, 1.0]) alpha = pm.Normal("alpha", 0, 1.5, size=3) pm.Normal("obs", alpha[index], np.sqrt(1e-2), observed=y) @@ -190,7 +190,7 @@ def test_shared_data_as_rv_input(self): See https://github.com/pymc-devs/pymc/issues/3842 """ with pm.Model() as m: - x = pm.MutableData("x", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) y = pm.Normal("y", mu=x, size=(2, 3)) assert y.eval().shape == (2, 3) idata = pm.sample( @@ -250,7 +250,7 @@ def test_shared_scalar_as_rv_input(self): def test_creation_of_data_outside_model_context(self): with pytest.raises((IndexError, TypeError)) as error: - pm.ConstantData("data", [1.1, 2.2, 3.3]) + pm.Data("data", [1.1, 2.2, 3.3]) error.match("No model on context stack") def test_set_data_to_non_data_container_variables(self): @@ -272,8 +272,8 @@ def test_set_data_to_non_data_container_variables(self): @pytest.mark.xfail(reason="Depends on ModelGraph") def test_model_to_graphviz_for_model_with_data_container(self, tmp_path): with pm.Model() as model: - x = pm.ConstantData("x", [1.0, 2.0, 3.0]) - y = pm.MutableData("y", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) + y = pm.Data("y", [1.0, 2.0, 3.0]) beta = pm.Normal("beta", 0, 10.0) obs_sigma = floatX(np.sqrt(1e-2)) pm.Normal("obs", beta * x, obs_sigma, observed=y) @@ -289,14 +289,14 @@ def test_model_to_graphviz_for_model_with_data_container(self, tmp_path): pm.model_to_graphviz(model, formatting=formatting) exp_without = [ - 'x [label="x\n~\nConstantData" shape=box style="rounded, filled"]', - 'y [label="x\n~\nMutableData" shape=box style="rounded, filled"]', + 'x [label="x\n~\\Data" shape=box style="rounded, filled"]', + 'y [label="x\n~\nData" shape=box style="rounded, filled"]', 'beta [label="beta\n~\nNormal"]', 'obs [label="obs\n~\nNormal" style=filled]', ] exp_with = [ - 'x [label="x\n~\nConstantData" shape=box style="rounded, filled"]', - 'y [label="x\n~\nMutableData" shape=box style="rounded, filled"]', + 'x [label="x\n~\nData" shape=box style="rounded, filled"]', + 'y [label="x\n~\nData" shape=box style="rounded, filled"]', 'beta [label="beta\n~\nNormal(mu=0.0, sigma=10.0)"]', f'obs [label="obs\n~\nNormal(mu=f(f(beta), x), sigma={obs_sigma})" style=filled]', ] @@ -324,10 +324,7 @@ def test_explicit_coords(self, seeded_test): } # pass coordinates explicitly, use numpy array in Data container with pm.Model(coords=coords) as pmodel: - # Dims created from coords are constant by default - assert isinstance(pmodel.dim_lengths["rows"], pt.TensorConstant) - assert isinstance(pmodel.dim_lengths["columns"], pt.TensorConstant) - pm.MutableData("observations", data, dims=("rows", "columns")) + pm.Data("observations", data, dims=("rows", "columns")) # new data with same (!) shape pm.set_data({"observations": data + 1}) # new data with same (!) shape and coords @@ -344,10 +341,8 @@ def test_explicit_coords(self, seeded_test): def test_set_coords_through_pmdata(self): with pm.Model() as pmodel: - pm.ConstantData( - "population", [100, 200], dims="city", coords={"city": ["Tinyvil", "Minitown"]} - ) - pm.MutableData( + pm.Data("population", [100, 200], dims="city", coords={"city": ["Tinyvil", "Minitown"]}) + pm.Data( "temperature", [[15, 20, 22, 17], [18, 22, 21, 12]], dims=("city", "season"), @@ -360,12 +355,12 @@ def test_set_coords_through_pmdata(self): def test_symbolic_coords(self): """ - In v4 dimensions can be created without passing coordinate values. + Since v4 dimensions can be created without passing coordinate values. Their lengths are then automatically linked to the corresponding Tensor dimension. """ with pm.Model() as pmodel: - # Dims created from MutableData are TensorVariables linked to the SharedVariable.shape - intensity = pm.MutableData("intensity", np.ones((2, 3)), dims=("row", "column")) + # Dims created from Data are TensorVariables linked to the SharedVariable.shape + intensity = pm.Data("intensity", np.ones((2, 3)), dims=("row", "column")) assert "row" in pmodel.dim_lengths assert "column" in pmodel.dim_lengths assert isinstance(pmodel.dim_lengths["row"], TensorVariable) @@ -385,7 +380,7 @@ def test_implicit_coords_series(self, seeded_test): name="sales", ) with pm.Model() as pmodel: - pm.ConstantData("sales", ser_sales, dims="date", export_index_as_coords=True) + pm.Data("sales", ser_sales, dims="date", infer_dims_and_coords=True) assert "date" in pmodel.coords assert len(pmodel.coords["date"]) == 22 @@ -403,9 +398,7 @@ def test_implicit_coords_dataframe(self, seeded_test): # infer coordinates from index and columns of the DataFrame with pm.Model() as pmodel: - pm.ConstantData( - "observations", df_data, dims=("rows", "columns"), export_index_as_coords=True - ) + pm.Data("observations", df_data, dims=("rows", "columns"), infer_dims_and_coords=True) assert "rows" in pmodel.coords assert "columns" in pmodel.coords @@ -415,8 +408,7 @@ def test_implicit_coords_xarray(self): xr = pytest.importorskip("xarray") data = xr.DataArray([[1, 2, 3], [4, 5, 6]], dims=("y", "x")) with pm.Model() as pmodel: - with pytest.warns(DeprecationWarning): - pm.ConstantData("observations", data, dims=("x", "y"), export_index_as_coords=True) + pm.Data("observations", data, dims=("x", "y"), infer_dims_and_coords=True) assert "x" in pmodel.coords assert "y" in pmodel.coords assert pmodel.named_vars_to_dims == {"observations": ("x", "y")} @@ -427,7 +419,7 @@ def test_data_kwargs(self): strict_value = True allow_downcast_value = False with pm.Model(): - data = pm.MutableData( + data = pm.Data( "mdata", value=[[1.0], [2.0], [3.0]], strict=strict_value, @@ -436,20 +428,13 @@ def test_data_kwargs(self): assert data.container.strict is strict_value assert data.container.allow_downcast is allow_downcast_value - def test_data_mutable_default_warning(self): - with pm.Model(): - with pytest.warns(UserWarning, match="`mutable` kwarg was not specified"): - data = pm.Data("x", [1, 2, 3]) - assert isinstance(data, pt.TensorConstant) - pass - def test_masked_array_error(self): with pm.Model(): with pytest.raises( NotImplementedError, match="Masked arrays or arrays with `nan` entries are not supported.", ): - pm.ConstantData("x", [0, 1, np.nan, 2]) + pm.Data("x", [0, 1, np.nan, 2]) def test_data_naming(): @@ -458,7 +443,7 @@ def test_data_naming(): not given model-relative names. """ with pm.Model("named_model") as model: - x = pm.ConstantData("x", [1.0, 2.0, 3.0]) + x = pm.Data("x", [1.0, 2.0, 3.0]) y = pm.Normal("y") assert y.name == "named_model::y" assert x.name == "named_model::x" diff --git a/tests/test_model_graph.py b/tests/test_model_graph.py index 963edb607a..fb95b5b3fb 100644 --- a/tests/test_model_graph.py +++ b/tests/test_model_graph.py @@ -102,9 +102,9 @@ def radon_model(): # Anonymous SharedVariables don't show up floor_measure = pytensor.shared(floor_measure) - floor_measure_offset = pm.MutableData("floor_measure_offset", 1) + floor_measure_offset = pm.Data("floor_measure_offset", 1) y_hat = a + b * floor_measure + floor_measure_offset - log_radon = pm.MutableData("log_radon", np.random.normal(1, 1, size=n_homes)) + log_radon = pm.Data("log_radon", np.random.normal(1, 1, size=n_homes)) y_like = pm.Normal("y_like", mu=y_hat, sigma=sigma_y, observed=log_radon) compute_graph = { @@ -163,13 +163,13 @@ def model_with_dims(): population = pm.HalfNormal("population", sigma=5, dims=("city")) - time = pm.ConstantData("time", [2014, 2015, 2016], dims="year") + time = pm.Data("time", [2014, 2015, 2016], dims="year") n = pm.Deterministic( "tax revenue", economics * population[None, :] * time[:, None], dims=("year", "city") ) - yobs = pm.MutableData("observed", np.ones((3, 4))) + yobs = pm.Data("observed", np.ones((3, 4))) L = pm.Normal("L", n, observed=yobs) compute_graph = { @@ -218,7 +218,7 @@ def model_observation_dtype_casting(): Model at the source of the following issue: https://github.com/pymc-devs/pymc/issues/5795 """ with pm.Model() as model: - data = pm.ConstantData("data", [0, 0, 1, 1], dtype=int) + data = pm.Data("data", np.array([0, 0, 1, 1], dtype=int)) p = pm.Beta("p", 1, 1) bern = pm.Bernoulli("response", p, observed=data) @@ -326,7 +326,7 @@ def model_with_different_descendants(): intermediate = pm.Deterministic("intermediate", a + b) pred = pm.Deterministic("pred", intermediate * 3) - obs = pm.ConstantData("obs", 1.75) + obs = pm.Data("obs", 1.75) L = pm.Normal("L", mu=1 + 0.5 * pred, observed=obs) diff --git a/tests/test_printing.py b/tests/test_printing.py index 26692ca569..95a1e812ec 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -199,10 +199,10 @@ def setup_class(self): import pymc as pm with pm.Model() as model: - a = pm.Normal("a", pm.MutableData("a_data", (2,))) - b = pm.Normal("b", pm.MutableData("b_data", (2, 3))) - c = pm.Normal("c", pm.ConstantData("c_data", (2,))) - d = pm.Normal("d", pm.ConstantData("d_data", (2, 3))) + a = pm.Normal("a", pm.Data("a_data", (2,))) + b = pm.Normal("b", pm.Data("b_data", (2, 3))) + c = pm.Normal("c", pm.Data("c_data", (2,))) + d = pm.Normal("d", pm.Data("d_data", (2, 3))) self.distributions = [a, b, c, d] # tuples of (formatting, include_params) @@ -212,7 +212,7 @@ def setup_class(self): r"a ~ Normal(2, 1)", r"b ~ Normal(, 1)", r"c ~ Normal(2, 1)", - r"d ~ Normal(, 1)", + r"d ~ Normal(, 1)", ], ("plain", False): [ r"a ~ Normal", @@ -224,7 +224,7 @@ def setup_class(self): r"$\text{a} \sim \operatorname{Normal}(2,~1)$", r"$\text{b} \sim \operatorname{Normal}(\text{},~1)$", r"$\text{c} \sim \operatorname{Normal}(2,~1)$", - r"$\text{d} \sim \operatorname{Normal}(\text{},~1)$", + r"$\text{d} \sim \operatorname{Normal}(\text{},~1)$", ], ("latex", False): [ r"$\text{a} \sim \operatorname{Normal}$", diff --git a/tests/variational/test_minibatch_rv.py b/tests/variational/test_minibatch_rv.py index 55e4bb73f7..10ab0914fc 100644 --- a/tests/variational/test_minibatch_rv.py +++ b/tests/variational/test_minibatch_rv.py @@ -163,8 +163,8 @@ def test_minibatch_parameter_and_value(self): total_size = 1000 with pm.Model(check_bounds=False) as m: - AD = pm.MutableData("AD", np.arange(total_size, dtype="float64")) - TD = pm.MutableData("TD", np.arange(total_size, dtype="float64")) + AD = pm.Data("AD", np.arange(total_size, dtype="float64")) + TD = pm.Data("TD", np.arange(total_size, dtype="float64")) minibatch_idx = minibatch_index(0, 10, size=(9,)) AD_mt = AD[minibatch_idx] From e35c40267e02bf339a6724aab48607dcd15e735a Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 4 Dec 2023 12:53:45 +0100 Subject: [PATCH 3/3] Implement model transform that freezes RV dims and Data --- .github/workflows/tests.yml | 1 + .../model/{conditioning.rst => transform.rst} | 11 ++- pymc/model/core.py | 37 ++++----- pymc/model/fgraph.py | 4 +- pymc/model/transform/optimization.py | 80 +++++++++++++++++++ tests/model/transform/test_optimization.py | 51 ++++++++++++ 6 files changed, 160 insertions(+), 24 deletions(-) rename docs/source/api/model/{conditioning.rst => transform.rst} (56%) create mode 100644 pymc/model/transform/optimization.py create mode 100644 tests/model/transform/test_optimization.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a735feccbb..b31fd0538e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -96,6 +96,7 @@ jobs: tests/model/test_fgraph.py tests/model/transform/test_basic.py tests/model/transform/test_conditioning.py + tests/model/transform/test_optimization.py tests/test_model_graph.py tests/ode/test_ode.py tests/ode/test_utils.py diff --git a/docs/source/api/model/conditioning.rst b/docs/source/api/model/transform.rst similarity index 56% rename from docs/source/api/model/conditioning.rst rename to docs/source/api/model/transform.rst index 8eae8d72ac..3e83176b17 100644 --- a/docs/source/api/model/conditioning.rst +++ b/docs/source/api/model/transform.rst @@ -5,7 +5,16 @@ Model Conditioning .. autosummary:: :toctree: generated/ - change_value_transforms do observe + change_value_transforms remove_value_transforms + + +Model Optimization +------------------ +.. currentmodule:: pymc.model.transform.optimization +.. autosummary:: + :toctree: generated/ + + freeze_dims_and_data diff --git a/pymc/model/core.py b/pymc/model/core.py index f117bb1b92..0ca9e57c26 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1121,8 +1121,7 @@ def set_data( for d, dname in enumerate(dims): length_tensor = self.dim_lengths[dname] - with pytensor.config.change_flags(cxx=""): - old_length = length_tensor.eval() + old_length = length_tensor.eval() new_length = values.shape[d] original_coords = self.coords.get(dname, None) new_coords = coords.get(dname, None) @@ -1404,24 +1403,22 @@ def create_value_var( else: transform = _default_transform(rv_var.owner.op, rv_var) - if value_var is not None: - if transform is not None: - raise ValueError("Cannot use transform when providing a pre-defined value_var") - elif transform is None: - # Create value variable with the same type as the RV - value_var = rv_var.type() - value_var.name = rv_var.name - if pytensor.config.compute_test_value != "off": - value_var.tag.test_value = rv_var.tag.test_value - else: - # Create value variable with the same type as the transformed RV - value_var = transform.forward(rv_var, *rv_var.owner.inputs).type() - value_var.name = f"{rv_var.name}_{transform.name}__" - value_var.tag.transform = transform - if pytensor.config.compute_test_value != "off": - value_var.tag.test_value = transform.forward( - rv_var, *rv_var.owner.inputs - ).tag.test_value + if value_var is None: + if transform is None: + # Create value variable with the same type as the RV + value_var = rv_var.type() + value_var.name = rv_var.name + if pytensor.config.compute_test_value != "off": + value_var.tag.test_value = rv_var.tag.test_value + else: + # Create value variable with the same type as the transformed RV + value_var = transform.forward(rv_var, *rv_var.owner.inputs).type() + value_var.name = f"{rv_var.name}_{transform.name}__" + value_var.tag.transform = transform + if pytensor.config.compute_test_value != "off": + value_var.tag.test_value = transform.forward( + rv_var, *rv_var.owner.inputs + ).tag.test_value _add_future_warning_tag(value_var) rv_var.tag.value_var = value_var diff --git a/pymc/model/fgraph.py b/pymc/model/fgraph.py index e2660f210d..49d7d9cbe4 100644 --- a/pymc/model/fgraph.py +++ b/pymc/model/fgraph.py @@ -321,9 +321,7 @@ def first_non_model_var(var): var, value, *dims = model_var.owner.inputs transform = model_var.owner.op.transform model.free_RVs.append(var) - # PyMC does not allow setting transform when we pass a value_var. Why? - model.create_value_var(var, transform=None, value_var=value) - model.rvs_to_transforms[var] = transform + model.create_value_var(var, transform=transform, value_var=value) model.set_initval(var, initval=None) elif isinstance(model_var.owner.op, ModelObservedRV): var, value, *dims = model_var.owner.inputs diff --git a/pymc/model/transform/optimization.py b/pymc/model/transform/optimization.py new file mode 100644 index 0000000000..651d22310c --- /dev/null +++ b/pymc/model/transform/optimization.py @@ -0,0 +1,80 @@ +# Copyright 2024 The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from pytensor import clone_replace +from pytensor.compile import SharedVariable +from pytensor.graph import FunctionGraph +from pytensor.tensor import constant + +from pymc import Model +from pymc.model.fgraph import ModelFreeRV, fgraph_from_model, model_from_fgraph + + +def freeze_dims_and_data(model: Model) -> Model: + """Recreate a Model with fixed RV dimensions and Data values. + + The dimensions of the pre-existing RVs will no longer follow changes to the coordinates. + Likewise, it will not be possible to update pre-existing Data in the new model. + + Note that any new RVs and Data created after calling this function will still be "unfrozen". + + This transformation may allow more performant sampling, or compiling model functions to backends that + are more restrictive about dynamic shapes such as JAX. + """ + fg, memo = fgraph_from_model(model) + + # Replace mutable dim lengths and data by constants + frozen_vars = { + memo[dim_length]: constant( + dim_length.get_value(), name=dim_length.name, dtype=dim_length.type.dtype + ) + for dim_length in model.dim_lengths.values() + if isinstance(dim_length, SharedVariable) + } + frozen_vars |= { + memo[data_var].owner.inputs[0]: constant( + data_var.get_value(), name=data_var.name, dtype=data_var.type.dtype + ) + for data_var in model.named_vars.values() + if isinstance(data_var, SharedVariable) + } + + old_outs, coords = fg.outputs, fg._coords # type: ignore + # Rebuild strict will force the recreation of RV nodes with updated static types + new_outs = clone_replace(old_outs, replace=frozen_vars, rebuild_strict=False) # type: ignore + for old_out, new_out in zip(old_outs, new_outs): + new_out.name = old_out.name + fg = FunctionGraph(outputs=new_outs, clone=False) + fg._coords = coords # type: ignore + + # Recreate value variables from new RVs to propagate static types to logp graphs + replacements = {} + for node in fg.apply_nodes: + if not isinstance(node.op, ModelFreeRV): + continue + rv, old_value, *dims = node.inputs + if dims is None: + continue + transform = node.op.transform + if transform is None: + new_value = rv.type() + else: + new_value = transform.forward(rv, *rv.owner.inputs).type() # type: ignore + new_value.name = old_value.name + replacements[old_value] = new_value + fg.replace_all(tuple(replacements.items()), import_missing=True) + + return model_from_fgraph(fg) + + +__all__ = ("freeze_dims_and_data",) diff --git a/tests/model/transform/test_optimization.py b/tests/model/transform/test_optimization.py new file mode 100644 index 0000000000..01e1d394f0 --- /dev/null +++ b/tests/model/transform/test_optimization.py @@ -0,0 +1,51 @@ +# Copyright 2024 The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from pytensor.graph import Constant + +from pymc.data import Data +from pymc.distributions import HalfNormal, Normal +from pymc.model import Model +from pymc.model.transform.optimization import freeze_dims_and_data + + +def test_freeze_existing_rv_dims_and_data(): + with Model(coords={"test_dim": range(5)}) as m: + std = Data("std", [1]) + x = HalfNormal("x", std, dims=("test_dim",)) + y = Normal("y", shape=x.shape[0] + 1) + + x_logp, y_logp = m.logp(sum=False) + + assert not isinstance(std, Constant) + assert x.type.shape == (None,) + assert y.type.shape == (None,) + assert x_logp.type.shape == (None,) + assert y_logp.type.shape == (None,) + + frozen_m = freeze_dims_and_data(m) + std, x, y = frozen_m["std"], frozen_m["x"], frozen_m["y"] + x_logp, y_logp = frozen_m.logp(sum=False) + assert isinstance(std, Constant) + assert x.type.shape == (5,) + assert y.type.shape == (6,) + assert x_logp.type.shape == (5,) + assert y_logp.type.shape == (6,) + + +def test_freeze_rv_dims_nothing_to_change(): + with Model(coords={"test_dim": range(5)}) as m: + x = HalfNormal("x", shape=(5,)) + y = Normal("y", shape=x.shape[0] + 1) + + assert m.point_logps() == freeze_dims_and_data(m).point_logps()