From 5a26eaf920b1bbf5d396c07239bd0c2affdd0d5f Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 10 Jun 2024 17:31:11 +0200 Subject: [PATCH 1/3] Remove deprecated Model functionality --- pymc/model/core.py | 25 ++----------------------- tests/model/test_core.py | 16 ---------------- 2 files changed, 2 insertions(+), 39 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 874d52f9bd6..5ad35498fa0 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -107,18 +107,10 @@ def __new__(cls, name, bases, dct, **kwargs): def __enter__(self): self.__class__.context_class.get_contexts().append(self) - # self._pytensor_config is set in Model.__new__ - self._config_context = None - if hasattr(self, "_pytensor_config"): - self._config_context = pytensor.config.change_flags(**self._pytensor_config) - self._config_context.__enter__() return self def __exit__(self, typ, value, traceback): self.__class__.context_class.get_contexts().pop() - # self._pytensor_config is set in Model.__new__ - if self._config_context: - self._config_context.__exit__(typ, value, traceback) dct[__enter__.__name__] = __enter__ dct[__exit__.__name__] = __exit__ @@ -485,13 +477,6 @@ def __new__(cls, *args, **kwargs): instance._parent = kwargs.get("model") else: instance._parent = cls.get_context(error_if_none=False) - pytensor_config = kwargs.get("pytensor_config", {}) - if pytensor_config: - warnings.warn( - "pytensor_config is deprecated. Use pytensor.config or pytensor.config.change_flags context manager instead.", - FutureWarning, - ) - instance._pytensor_config = pytensor_config return instance @staticmethod @@ -507,10 +492,9 @@ def __init__( check_bounds=True, *, coords_mutable=None, - pytensor_config=None, model=None, ): - del pytensor_config, model # used in __new__ + del model # used in __new__ self.name = self._validate_name(name) self.check_bounds = check_bounds @@ -560,11 +544,6 @@ def __init__( functools.partial(str_for_model, formatting="latex"), self ) - @property - def model(self): - warnings.warn("Model.model property is deprecated. Just use Model.", FutureWarning) - return self - @property def parent(self): return self._parent @@ -671,7 +650,7 @@ def compile_d2logp( jacobian : bool Whether to include jacobian terms in logprob graph. Defaults to True. """ - return self.model.compile_fn( + return self.compile_fn( self.d2logp(vars=vars, jacobian=jacobian, negate_output=negate_output), **compile_kwargs, ) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index c00250b7399..fe4932b5722 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -1077,22 +1077,6 @@ def test_compile_fn(): np.testing.assert_allclose(result_compute, result_expect) -def test_model_pytensor_config(): - assert pytensor.config.mode != "JAX" - with pytest.warns(FutureWarning, match="pytensor_config is deprecated"): - m = pm.Model(pytensor_config=dict(mode="JAX")) - with m: - assert pytensor.config.mode == "JAX" - assert pytensor.config.mode != "JAX" - - -def test_deprecated_model_property(): - m = pm.Model() - with pytest.warns(FutureWarning, match="Model.model property is deprecated"): - m_property = m.model - assert m is m_property - - def test_model_parent_set_programmatically(): with pm.Model() as model: x = pm.Normal("x") From b77747d7a41a5df767f7bdfe40622f26ccdb0901 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 10 Jun 2024 17:33:44 +0200 Subject: [PATCH 2/3] More relevant documentation examples for Model --- pymc/model/core.py | 143 +++++++++++++++++++++++---------------- tests/model/test_core.py | 26 +++++++ 2 files changed, 111 insertions(+), 58 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 5ad35498fa0..7aec544d79f 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -392,76 +392,103 @@ class Model(WithMemoization, metaclass=ContextMeta): name : str name that will be used as prefix for names of all random variables defined within model + coords : dict + Xarray-like coordinate keys and values. These coordinates can be used + to specify the shape of random variables and to label (but not specify) + the shape of Determinsitic, Potential and Data objects. + Other than specifying the shape of random variables, coordinates have no + effect on the model. They can't be used for label-based broadcasting or indexing. + You must use numpy-like operations for those behaviors. check_bounds : bool Ensure that input parameters to distributions are in a valid range. If your model is built in a way where you know your parameters can only take on valid values you can set this to False for increased speed. This should not be used if your model contains discrete variables. + model : PyMC model, optional + A parent model that this model belongs to. If not specified and the current model + is created inside another model's context, the parent model will be set to that model. + If `None` the model will not have a parent. Examples -------- - How to define a custom model + Use context manager to define model and respective variables .. code-block:: python - class CustomModel(Model): - # 1) override init - def __init__(self, mean=0, sigma=1, name=''): - # 2) call super's init first, passing model and name - # to it name will be prefix for all variables here if - # no name specified for model there will be no prefix - super().__init__(name, model) - # now you are in the context of instance, - # `modelcontext` will return self you can define - # variables in several ways note, that all variables - # will get model's name prefix - - # 3) you can create variables with the register_rv method - self.register_rv(Normal.dist(mu=mean, sigma=sigma), 'v1', initval=1) - # this will create variable named like '{name::}v1' - # and assign attribute 'v1' to instance created - # variable can be accessed with self.v1 or self['v1'] - - # 4) this syntax will also work as we are in the - # context of instance itself, names are given as usual - Normal('v2', mu=mean, sigma=sigma) - - # something more complex is allowed, too - half_cauchy = HalfCauchy('sigma', beta=10, initval=1.) - Normal('v3', mu=mean, sigma=half_cauchy) - - # Deterministic variables can be used in usual way - Deterministic('v3_sq', self.v3 ** 2) - - # Potentials too - Potential('p1', pt.constant(1)) - - # After defining a class CustomModel you can use it in several - # ways - - # I: - # state the model within a context - with Model() as model: - CustomModel() - # arbitrary actions - - # II: - # use new class as entering point in context - with CustomModel() as model: - Normal('new_normal_var', mu=1, sigma=0) - - # III: - # just get model instance with all that was defined in it - model = CustomModel() - - # IV: - # use many custom models within one context - with Model() as model: - CustomModel(mean=1, name='first') - CustomModel(mean=2, name='second') - - # variables inside both scopes will be named like `first::*`, `second::*` + import pymc as pm + + with pm.Model() as model: + x = pm.Normal("x") + + + Use object API to define model and respective variables + + .. code-block:: python + + import pymc as pm + + model = pm.Model() + x = pm.Normal("x", model=model) + + + Use coords for defining the shape of random variables and labeling other model variables + + .. code-block:: python + + import pymc as pm + import numpy as np + + coords = { + "feature", ["A", "B", "C"], + "trial", [1, 2, 3, 4, 5], + } + + with pm.Model(coords=coords) as model: + intercept = pm.Normal("intercept", shape=(3,)) # Variable will have default dim label `intercept__dim_0` + beta = pm.Normal("beta", dims=("feature",)) # Variable will have shape (3,) and dim label `feature` + + # Dims below are only used for labeling, they have no effect on shape + idx = pm.Data("idx", np.array([0, 1, 1, 2, 2])) # Variable will have default dim label `idx__dim_0` + x = pm.Data("x", np.random.normal(size=(5, 3)), dims=("trial", "feature")) + mu = pm.Deterministic("mu", intercept[idx] + beta @ x, dims="trial") # single dim can be passed as string + + # Dims controls the shape of the variable + # If not specified, it would be inferred from the shape of the observations + y = pm.Normal("y", mu=mu, observed=[-1, 0, 0, 1, 1], dims=("trial",)) + + + Define nested models, and provide name for variable name prefixing + + .. code-block:: python + + import pymc as pm + + with pm.Model(name="root") as root: + x = pm.Normal("x") # Variable wil be named "root::x" + + with pm.Model(name='first') as first: + # Variable will belong to root and first + y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y" + + # Can pass parent model explicitly + with pm.Model(name='second', model=root) as second: + # Variable will belong to root and second + z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" + + + Set `check_bounds` to False for models with only continuous variables and default transformers + PyMC will remove the bounds check from the model logp which can speed up sampling + + .. code-block:: python + + import pymc as pm + + with pm.Model(check_bounds=False) as model: + sigma = pm.HalfNormal("sigma") + x = pm.Normal("x", sigma=sigma) # No bounds check will be performed on `sigma` + + """ if TYPE_CHECKING: diff --git a/tests/model/test_core.py b/tests/model/test_core.py index fe4932b5722..484ac76d9c6 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -130,6 +130,27 @@ def test_context_passes_vars_to_parent_model(self): assert m["d"] is model["one_more::d"] assert m["one_more::d"] is model["one_more::d"] + def test_docstring_example(self): + with pm.Model(name="root") as root: + x = pm.Normal("x") # Variable wil be named "root::x" + + with pm.Model(name="first") as first: + # Variable will belong to root and first + y = pm.Normal("y", mu=x) # Variable wil be named "root::first::y" + + # Can pass parent model explicitly + with pm.Model(name="second", model=root) as second: + # Variable will belong to root and second + z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" + + assert x.name == "root::x" + assert y.name == "root::first::y" + assert z.name == "root::second::z" + + assert set(root.basic_RVs) == {x, y, z} + assert set(first.basic_RVs) == {y} + assert set(second.basic_RVs) == {z} + class TestNested: def test_nest_context_works(self): @@ -1084,7 +1105,12 @@ def test_model_parent_set_programmatically(): with pm.Model(model=model): y = pm.Normal("y") + with model: + with pm.Model(model=None): + z = pm.Normal("z") + assert "y" in model.named_vars + assert "z" in model.named_vars class TestModelContext: From 269da5c03c771a54fc0478704a4a3bb53d4aa9c8 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Mon, 10 Jun 2024 18:14:02 +0200 Subject: [PATCH 3/3] Allow opting out of model nesting --- pymc/model/core.py | 20 +++++++++----- pymc/model/fgraph.py | 4 +-- pymc/model/transform/basic.py | 2 +- pymc/model/transform/conditioning.py | 2 +- pymc/sampling/deterministic.py | 2 +- pymc/stats/log_density.py | 39 +++++++++++----------------- tests/model/test_core.py | 17 ++++++++++-- tests/model/test_fgraph.py | 11 +++++--- 8 files changed, 55 insertions(+), 42 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 7aec544d79f..475ff42b5f7 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -25,6 +25,7 @@ Literal, Optional, TypeVar, + Union, cast, overload, ) @@ -441,7 +442,7 @@ class Model(WithMemoization, metaclass=ContextMeta): coords = { "feature", ["A", "B", "C"], - "trial", [1, 2, 3, 4, 5], + "trial", [1, 2, 3, 4, 5], } with pm.Model(coords=coords) as model: @@ -476,6 +477,11 @@ class Model(WithMemoization, metaclass=ContextMeta): # Variable will belong to root and second z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" + # Set None for standalone model + with pm.Model(name="third", model=None) as third: + # Variable will belong to third only + w = pm.Normal("w") # Variable wil be named "third::w" + Set `check_bounds` to False for models with only continuous variables and default transformers PyMC will remove the bounds check from the model logp which can speed up sampling @@ -497,13 +503,13 @@ def __enter__(self: Self) -> Self: ... def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None: ... - def __new__(cls, *args, **kwargs): + def __new__(cls, *args, model: Union[Literal[UNSET], None, "Model"] = UNSET, **kwargs): # resolves the parent instance instance = super().__new__(cls) - if kwargs.get("model") is not None: - instance._parent = kwargs.get("model") - else: + if model is UNSET: instance._parent = cls.get_context(error_if_none=False) + else: + instance._parent = model return instance @staticmethod @@ -519,9 +525,9 @@ def __init__( check_bounds=True, *, coords_mutable=None, - model=None, + model: Union[Literal[UNSET], None, "Model"] = UNSET, ): - del model # used in __new__ + del model # used in __new__ to define the parent of this model self.name = self._validate_name(name) self.check_bounds = check_bounds diff --git a/pymc/model/fgraph.py b/pymc/model/fgraph.py index ce15c40760f..b1d67fd07b0 100644 --- a/pymc/model/fgraph.py +++ b/pymc/model/fgraph.py @@ -299,9 +299,7 @@ def first_non_model_var(var): else: return var - model = Model() - if model.parent is not None: - raise RuntimeError("model_to_fgraph cannot be called inside a PyMC model context") + model = Model(model=None) # Do not inherit from any model in the context manager _coords = getattr(fgraph, "_coords", {}) _dim_lengths = getattr(fgraph, "_dim_lengths", {}) diff --git a/pymc/model/transform/basic.py b/pymc/model/transform/basic.py index aff6042de5a..76556ae08ab 100644 --- a/pymc/model/transform/basic.py +++ b/pymc/model/transform/basic.py @@ -16,7 +16,7 @@ from pytensor import Variable from pytensor.graph import ancestors -from pymc import Model +from pymc.model.core import Model from pymc.model.fgraph import ( ModelObservedRV, ModelVar, diff --git a/pymc/model/transform/conditioning.py b/pymc/model/transform/conditioning.py index 0979964eecf..531ec2bd6b9 100644 --- a/pymc/model/transform/conditioning.py +++ b/pymc/model/transform/conditioning.py @@ -19,9 +19,9 @@ from pytensor.graph import ancestors from pytensor.tensor import TensorVariable -from pymc import Model from pymc.logprob.transforms import Transform from pymc.logprob.utils import rvs_in_graph +from pymc.model.core import Model from pymc.model.fgraph import ( ModelDeterministic, ModelFreeRV, diff --git a/pymc/sampling/deterministic.py b/pymc/sampling/deterministic.py index b0b04f38ec7..3d8398c3a7e 100644 --- a/pymc/sampling/deterministic.py +++ b/pymc/sampling/deterministic.py @@ -83,7 +83,7 @@ def compute_deterministics( model = modelcontext(model) if var_names is None: - deterministics = model.deterministics + deterministics = list(model.deterministics) var_names = [det.name for det in deterministics] else: deterministics = [model[var_name] for var_name in var_names] diff --git a/pymc/stats/log_density.py b/pymc/stats/log_density.py index 5b6406d02b1..a26f8aa60df 100644 --- a/pymc/stats/log_density.py +++ b/pymc/stats/log_density.py @@ -25,6 +25,8 @@ __all__ = ("compute_log_likelihood", "compute_log_prior") +from pymc.model.transform.conditioning import remove_value_transforms + def compute_log_likelihood( idata: InferenceData, @@ -126,46 +128,35 @@ def compute_log_density( if kind not in ("likelihood", "prior"): raise ValueError("kind must be either 'likelihood' or 'prior'") + # We need to disable transforms, because the InferenceData only keeps the untransformed values + umodel = remove_value_transforms(model) + if kind == "likelihood": - target_rvs = model.observed_RVs + target_rvs = list(umodel.observed_RVs) target_str = "observed_RVs" else: - target_rvs = model.free_RVs + target_rvs = list(umodel.free_RVs) target_str = "free_RVs" if var_names is None: vars = target_rvs var_names = tuple(rv.name for rv in vars) else: - vars = [model.named_vars[name] for name in var_names] + vars = [umodel.named_vars[name] for name in var_names] if not set(vars).issubset(target_rvs): raise ValueError(f"var_names must refer to {target_str} in the model. Got: {var_names}") - # We need to temporarily disable transforms, because the InferenceData only keeps the untransformed values - try: - original_rvs_to_values = model.rvs_to_values - original_rvs_to_transforms = model.rvs_to_transforms - - model.rvs_to_values = { - rv: rv.clone() if rv not in model.observed_RVs else value - for rv, value in model.rvs_to_values.items() - } - model.rvs_to_transforms = {rv: None for rv in model.basic_RVs} - - elemwise_logdens_fn = model.compile_fn( - inputs=model.value_vars, - outs=model.logp(vars=vars, sum=False), - on_unused_input="ignore", - ) - finally: - model.rvs_to_values = original_rvs_to_values - model.rvs_to_transforms = original_rvs_to_transforms + elemwise_logdens_fn = umodel.compile_fn( + inputs=umodel.value_vars, + outs=umodel.logp(vars=vars, sum=False), + on_unused_input="ignore", + ) - coords, dims = coords_and_dims_for_inferencedata(model) + coords, dims = coords_and_dims_for_inferencedata(umodel) logdens_dataset = apply_function_over_dataset( elemwise_logdens_fn, - posterior[[rv.name for rv in model.free_RVs]], + posterior[[rv.name for rv in umodel.free_RVs]], output_var_names=var_names, sample_dims=sample_dims, dims=dims, diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 484ac76d9c6..e43a86cebfe 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -143,13 +143,20 @@ def test_docstring_example(self): # Variable will belong to root and second z = pm.Normal("z", mu=y) # Variable wil be named "root::second::z" + # Set None for standalone model + with pm.Model(name="third", model=None) as third: + # Variable will belong to third only + w = pm.Normal("w") # Variable wil be named "third::w" + assert x.name == "root::x" assert y.name == "root::first::y" assert z.name == "root::second::z" + assert w.name == "third::w" assert set(root.basic_RVs) == {x, y, z} assert set(first.basic_RVs) == {y} assert set(second.basic_RVs) == {z} + assert set(third.basic_RVs) == {w} class TestNested: @@ -1106,11 +1113,17 @@ def test_model_parent_set_programmatically(): y = pm.Normal("y") with model: + # Default inherits from model + with pm.Model(): + z_in = pm.Normal("z_in") + + # Explict None opts out of model context with pm.Model(model=None): - z = pm.Normal("z") + z_out = pm.Normal("z_out") assert "y" in model.named_vars - assert "z" in model.named_vars + assert "z_in" in model.named_vars + assert "z_out" not in model.named_vars class TestModelContext: diff --git a/tests/model/test_fgraph.py b/tests/model/test_fgraph.py index 7a57bfc16a4..a964f1faf6d 100644 --- a/tests/model/test_fgraph.py +++ b/tests/model/test_fgraph.py @@ -267,10 +267,15 @@ def test_context_error(): with pm.Model() as m: x = pm.Normal("x") - fg = fgraph_from_model(m) + fg, _ = fgraph_from_model(m) - with pytest.raises(RuntimeError, match="cannot be called inside a PyMC model context"): - model_from_fgraph(fg) + new_m = model_from_fgraph(fg) + new_x = new_m["x"] + + assert new_m.parent is None + assert x != new_x + assert m.named_vars == {"x": x} + assert new_m.named_vars == {"x": new_x} def test_sub_model_error():