From 8a50f9f7a2d4b7b4ef3000cd8019ec765676b45a Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Thu, 21 Mar 2024 21:29:26 +0300 Subject: [PATCH 01/24] add default transform argument --- pymc/distributions/distribution.py | 2 ++ pymc/model/core.py | 44 ++++++++++++++++++++++++------ 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 6f48673dc9..5821bd354d 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -333,6 +333,7 @@ def __new__( observed=None, total_size=None, transform=UNSET, + default_transform=UNSET, **kwargs, ) -> TensorVariable: """Adds a tensor variable corresponding to a PyMC distribution to the current model. @@ -418,6 +419,7 @@ def __new__( total_size, dims=dims, transform=transform, + default_transform=default_transform, initval=initval, ) diff --git a/pymc/model/core.py b/pymc/model/core.py index eee8f2a904..67625426a1 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -48,7 +48,7 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import GenTensorVariable, is_minibatch -from pymc.distributions.transforms import _default_transform +from pymc.distributions.transforms import ChainedTransform, _default_transform from pymc.exceptions import ( BlockModelAccessError, ImputationWarning, @@ -1214,7 +1214,15 @@ def set_data( shared_object.set_value(values) def register_rv( - self, rv_var, name, observed=None, total_size=None, dims=None, transform=UNSET, initval=None + self, + rv_var, + name, + observed=None, + total_size=None, + dims=None, + transform=UNSET, + default_transform=UNSET, + initval=None, ): """Register an (un)observed random variable with the model. @@ -1255,7 +1263,7 @@ def register_rv( if total_size is not None: raise ValueError("total_size can only be passed to observed RVs") self.free_RVs.append(rv_var) - self.create_value_var(rv_var, transform) + self.create_value_var(rv_var, transform, default_transform) self.add_named_variable(rv_var, dims) self.set_initval(rv_var, initval) else: @@ -1278,7 +1286,9 @@ def register_rv( # `rv_var` is potentially changed by `make_obs_var`, # for example into a new graph for imputation of missing data. - rv_var = self.make_obs_var(rv_var, observed, dims, transform, total_size) + rv_var = self.make_obs_var( + rv_var, observed, dims, transform, default_transform, total_size + ) return rv_var @@ -1288,6 +1298,7 @@ def make_obs_var( data: np.ndarray, dims, transform: Any | None, + default_transform: Any | None, total_size: int | None, ) -> TensorVariable: """Create a `TensorVariable` for an observed random variable. @@ -1344,7 +1355,12 @@ def make_obs_var( self.observed_RVs.append(observed_rv) # Register FreeRV corresponding to unobserved components - self.register_rv(unobserved_rv, f"{name}_unobserved", transform=transform) + self.register_rv( + unobserved_rv, + f"{name}_unobserved", + transform=transform, + default_transform=default_transform, + ) # Register Deterministic that combines observed and missing # Note: This can widely increase memory consumption during sampling for large datasets @@ -1370,7 +1386,11 @@ def make_obs_var( return rv_var def create_value_var( - self, rv_var: TensorVariable, transform: Any, value_var: Variable | None = None + self, + rv_var: TensorVariable, + transform: Any, + default_transform: Any, + value_var: Variable | None = None, ) -> TensorVariable: """Create a ``TensorVariable`` that will be used as the random variable's "value" in log-likelihood graphs. @@ -1396,11 +1416,17 @@ def create_value_var( # Make the value variable a transformed value variable, # if there's an applicable transform - if transform is UNSET: + + if default_transform is UNSET: if rv_var.owner is None: - transform = None + default_transform = None else: - transform = _default_transform(rv_var.owner.op, rv_var) + default_transform = _default_transform(rv_var.owner.op, rv_var) + + if transform is UNSET: + transform = default_transform + else: + transform = ChainedTransform([default_transform, transform]) if value_var is None: if transform is None: From 7434b9a78887beb9bf60b1943465d17974bcfb4a Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Fri, 22 Mar 2024 13:36:12 +0300 Subject: [PATCH 02/24] Apply suggestions from code review Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- pymc/model/core.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pymc/model/core.py b/pymc/model/core.py index 67625426a1..1e28c7355c 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1217,6 +1217,7 @@ def register_rv( self, rv_var, name, + *, observed=None, total_size=None, dims=None, @@ -1388,6 +1389,7 @@ def make_obs_var( def create_value_var( self, rv_var: TensorVariable, + *, transform: Any, default_transform: Any, value_var: Variable | None = None, From 4fde32f6bde8abab54e6094a04a41b745c8ab506 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Fri, 22 Mar 2024 14:36:12 +0300 Subject: [PATCH 03/24] fix arguments passing --- pymc/distributions/distribution.py | 4 ++-- pymc/model/core.py | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 5821bd354d..3af0bb110e 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -415,8 +415,8 @@ def __new__( rv_out = model.register_rv( rv_out, name, - observed, - total_size, + observed=observed, + total_size=total_size, dims=dims, transform=transform, default_transform=default_transform, diff --git a/pymc/model/core.py b/pymc/model/core.py index 1e28c7355c..c36ce6fe91 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1264,7 +1264,7 @@ def register_rv( if total_size is not None: raise ValueError("total_size can only be passed to observed RVs") self.free_RVs.append(rv_var) - self.create_value_var(rv_var, transform, default_transform) + self.create_value_var(rv_var, transform=transform, default_transform=default_transform) self.add_named_variable(rv_var, dims) self.set_initval(rv_var, initval) else: @@ -1380,7 +1380,9 @@ def make_obs_var( rv_var.name = name rv_var.tag.observations = data - self.create_value_var(rv_var, transform=None, value_var=data) + self.create_value_var( + rv_var, transform=transform, default_transform=default_transform, value_var=data + ) self.add_named_variable(rv_var, dims) self.observed_RVs.append(rv_var) @@ -1418,7 +1420,6 @@ def create_value_var( # Make the value variable a transformed value variable, # if there's an applicable transform - if default_transform is UNSET: if rv_var.owner is None: default_transform = None From 5c7519f3420d3fe54bb3b6261f887974ad02bcf2 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Fri, 22 Mar 2024 16:18:16 +0300 Subject: [PATCH 04/24] add support for transform=None --- pymc/distributions/distribution.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 3af0bb110e..b6a2eb17ec 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -398,6 +398,16 @@ def __new__( if not isinstance(name, string_types): raise TypeError(f"Name needs to be a string but got: {name}") + if transform is None: + transform = UNSET + default_transform = None + warnings.warn( + "To disable default transform, please use default_transform=None" + " instead of transform=None. Setting transform to None will" + " not take effect in future.", + UserWarning, + ) + dims = convert_dims(dims) if observed is not None: observed = convert_observed_data(observed) From 1dc0fcb3b98ca74447c01822e57c50fd9b08a990 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Fri, 22 Mar 2024 17:37:31 +0300 Subject: [PATCH 05/24] fix failed tests --- pymc/model/core.py | 6 ++++-- tests/distributions/test_mixture.py | 10 +++++++++- tests/distributions/test_transform.py | 4 ++-- tests/sampling/test_mcmc.py | 2 +- 4 files changed, 16 insertions(+), 6 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index c36ce6fe91..e5a03f56d2 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1351,7 +1351,9 @@ def make_obs_var( # Register ObservedRV corresponding to observed component observed_rv.name = f"{name}_observed" - self.create_value_var(observed_rv, transform=None, value_var=observed_data) + self.create_value_var( + observed_rv, transform=transform, default_transform=None, value_var=observed_data + ) self.add_named_variable(observed_rv) self.observed_RVs.append(observed_rv) @@ -1428,7 +1430,7 @@ def create_value_var( if transform is UNSET: transform = default_transform - else: + elif transform and default_transform: transform = ChainedTransform([default_transform, transform]) if value_var is None: diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index a7293c0ed3..2ec9aaec49 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -1359,7 +1359,15 @@ def test_warning(self): with warnings.catch_warnings(): warnings.simplefilter("error") - Mixture("mix4", w=[0.5, 0.5], comp_dists=comp_dists, transform=None) + Mixture("mix4", w=[0.5, 0.5], comp_dists=comp_dists, default_transform=None) + + with pytest.warns( + UserWarning, + match="To disable default transform, please use default_transform=None" + " instead of transform=None. Setting transform to None will" + " not take effect in future.", + ): + Mixture("mix5", w=[0.5, 0.5], comp_dists=comp_dists, transform=None) with warnings.catch_warnings(): warnings.simplefilter("error") diff --git a/tests/distributions/test_transform.py b/tests/distributions/test_transform.py index b0187a4ebe..4bc5e8543d 100644 --- a/tests/distributions/test_transform.py +++ b/tests/distributions/test_transform.py @@ -619,7 +619,7 @@ def test_transform_univariate_dist_logp_shape(): def test_univariate_transform_multivariate_dist_raises(): with pm.Model() as m: - pm.Dirichlet("x", [1, 1, 1], transform=tr.log) + pm.Dirichlet("x", [1, 1, 1], transform=tr.log, default_transform=None) for jacobian_val in (True, False): with pytest.raises( @@ -645,7 +645,7 @@ def log_jac_det(self, value, *inputs): buggy_transform = BuggyTransform() with pm.Model() as m: - pm.Uniform("x", shape=(4, 3), transform=buggy_transform) + pm.Uniform("x", shape=(4, 3), transform=buggy_transform, default_transform=None) for jacobian_val in (True, False): with pytest.raises( diff --git a/tests/sampling/test_mcmc.py b/tests/sampling/test_mcmc.py index 3f676d0846..31b48250fc 100644 --- a/tests/sampling/test_mcmc.py +++ b/tests/sampling/test_mcmc.py @@ -303,7 +303,7 @@ def test_transform_with_rv_dependency(self, symbolic_rv): transform = pm.distributions.transforms.Interval( bounds_fn=lambda *inputs: (inputs[-2], inputs[-1]) ) - y = pm.Uniform("y", lower=0, upper=x, transform=transform) + y = pm.Uniform("y", lower=0, upper=x, transform=transform, default_transform=None) with warnings.catch_warnings(): warnings.filterwarnings("ignore", ".*number of samples.*", UserWarning) trace = pm.sample(tune=10, draws=50, return_inferencedata=False, random_seed=336) From 55b0edb0fe0fa7140cdd7d8718fbe4bf125fb0f9 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Tue, 26 Mar 2024 17:02:27 +0300 Subject: [PATCH 06/24] Update pymc/distributions/distribution.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- pymc/distributions/distribution.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index b6a2eb17ec..36711383dd 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -398,13 +398,12 @@ def __new__( if not isinstance(name, string_types): raise TypeError(f"Name needs to be a string but got: {name}") - if transform is None: - transform = UNSET + if transform is None and default_transform is UNSET: default_transform = None warnings.warn( "To disable default transform, please use default_transform=None" " instead of transform=None. Setting transform to None will" - " not take effect in future.", + " not have any effect in future.", UserWarning, ) From d791ac7241168f7da147fca255995375cbf3d1bd Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 26 Mar 2024 17:47:13 +0300 Subject: [PATCH 07/24] fix failed tests --- pymc/model/core.py | 2 +- pymc/model/fgraph.py | 4 ++-- tests/distributions/test_mixture.py | 8 ++++---- tests/model/test_core.py | 10 +++++----- tests/model/transform/test_conditioning.py | 4 ++-- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index e5a03f56d2..db019b8008 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1383,7 +1383,7 @@ def make_obs_var( rv_var.tag.observations = data self.create_value_var( - rv_var, transform=transform, default_transform=default_transform, value_var=data + rv_var, transform=transform, default_transform=None, value_var=data ) self.add_named_variable(rv_var, dims) self.observed_RVs.append(rv_var) diff --git a/pymc/model/fgraph.py b/pymc/model/fgraph.py index 05d8fe4200..b71c168f08 100644 --- a/pymc/model/fgraph.py +++ b/pymc/model/fgraph.py @@ -320,12 +320,12 @@ def first_non_model_var(var): var, value, *dims = model_var.owner.inputs transform = model_var.owner.op.transform model.free_RVs.append(var) - model.create_value_var(var, transform=transform, value_var=value) + model.create_value_var(var, transform=transform, default_transform=None, value_var=value) model.set_initval(var, initval=None) elif isinstance(model_var.owner.op, ModelObservedRV): var, value, *dims = model_var.owner.inputs model.observed_RVs.append(var) - model.create_value_var(var, transform=None, value_var=value) + model.create_value_var(var, transform=None, default_transform=None, value_var=value) elif isinstance(model_var.owner.op, ModelPotential): var, *dims = model_var.owner.inputs model.potentials.append(var) diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index 2ec9aaec49..97122ea6f8 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -1364,20 +1364,20 @@ def test_warning(self): with pytest.warns( UserWarning, match="To disable default transform, please use default_transform=None" - " instead of transform=None. Setting transform to None will" - " not take effect in future.", + " instead of transform=None. Setting transform to None will not have" + " any effect in future.", ): Mixture("mix5", w=[0.5, 0.5], comp_dists=comp_dists, transform=None) with warnings.catch_warnings(): warnings.simplefilter("error") - Mixture("mix5", w=[0.5, 0.5], comp_dists=comp_dists, observed=1) + Mixture("mix6", w=[0.5, 0.5], comp_dists=comp_dists, observed=1) # Case where the appropriate default transform is None comp_dists = [Normal.dist(), Normal.dist()] with warnings.catch_warnings(): warnings.simplefilter("error") - Mixture("mix6", w=[0.5, 0.5], comp_dists=comp_dists) + Mixture("mix7", w=[0.5, 0.5], comp_dists=comp_dists) class TestZeroInflatedMixture: diff --git a/tests/model/test_core.py b/tests/model/test_core.py index d62890ab1d..c5eaec430e 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -549,18 +549,18 @@ def test_make_obs_var(): # The function requires data and RV dimensionality to be compatible with pytest.raises(ShapeError, match="Dimensionality of data and RV don't match."): - fake_model.make_obs_var(fake_distribution, np.ones((3, 3, 1)), None, None, None) + fake_model.make_obs_var(fake_distribution, np.ones((3, 3, 1)), None, None, None, None) # Check function behavior using the various inputs # dense, sparse: Ensure that the missing values are appropriately set to None # masked: a deterministic variable is returned - dense_output = fake_model.make_obs_var(fake_distribution, dense_input, None, None, None) + dense_output = fake_model.make_obs_var(fake_distribution, dense_input, None, None, None, None) assert dense_output == fake_distribution assert isinstance(fake_model.rvs_to_values[dense_output], TensorConstant) del fake_model.named_vars[fake_distribution.name] - sparse_output = fake_model.make_obs_var(fake_distribution, sparse_input, None, None, None) + sparse_output = fake_model.make_obs_var(fake_distribution, sparse_input, None, None, None, None) assert sparse_output == fake_distribution assert sparse.basic._is_sparse_variable(fake_model.rvs_to_values[sparse_output]) del fake_model.named_vars[fake_distribution.name] @@ -568,7 +568,7 @@ def test_make_obs_var(): # Here the RandomVariable is split into observed/imputed and a Deterministic is returned with pytest.warns(ImputationWarning): masked_output = fake_model.make_obs_var( - fake_distribution, masked_array_input, None, None, None + fake_distribution, masked_array_input, None, None, None, None ) assert masked_output != fake_distribution assert not isinstance(masked_output, RandomVariable) @@ -581,7 +581,7 @@ def test_make_obs_var(): # Test that setting total_size returns a MinibatchRandomVariable scaled_outputs = fake_model.make_obs_var( - fake_distribution, dense_input, None, None, total_size=100 + fake_distribution, dense_input, None, None, None, total_size=100 ) assert scaled_outputs != fake_distribution assert isinstance(scaled_outputs.owner.op, MinibatchRandomVariable) diff --git a/tests/model/transform/test_conditioning.py b/tests/model/transform/test_conditioning.py index a9a8ab712d..25f4693052 100644 --- a/tests/model/transform/test_conditioning.py +++ b/tests/model/transform/test_conditioning.py @@ -286,8 +286,8 @@ def test_change_value_transforms_error(): def test_remove_value_transforms(): with pm.Model() as base_m: - p = pm.Uniform("p", transform=logodds) - q = pm.Uniform("q", transform=logodds) + p = pm.Uniform("p", transform=logodds, default_transform=None) + q = pm.Uniform("q", transform=logodds, default_transform=None) new_m = remove_value_transforms(base_m) new_p = new_m["p"] From def0f6ee4fed314b7b942ad7878aad3a5345407c Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 26 Mar 2024 18:18:19 +0300 Subject: [PATCH 08/24] fix test --- tests/logprob/test_utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/logprob/test_utils.py b/tests/logprob/test_utils.py index c59b332495..7d3094c5b5 100644 --- a/tests/logprob/test_utils.py +++ b/tests/logprob/test_utils.py @@ -218,11 +218,15 @@ def test_interdependent_transformed_rvs(self, reversed): transform = pm.distributions.transforms.Interval( bounds_fn=lambda *inputs: (inputs[-2], inputs[-1]) ) - x = pm.Uniform("x", lower=0, upper=1, transform=transform) + x = pm.Uniform("x", lower=0, upper=1, transform=transform, default_transform=None) # Operation between the variables provides a regression test for #7054 - y = pm.Uniform("y", lower=0, upper=pt.exp(x), transform=transform) - z = pm.Uniform("z", lower=0, upper=y, transform=transform) - w = pm.Uniform("w", lower=0, upper=pt.square(z), transform=transform) + y = pm.Uniform( + "y", lower=0, upper=pt.exp(x), transform=transform, default_transform=None + ) + z = pm.Uniform("z", lower=0, upper=y, transform=transform, default_transform=None) + w = pm.Uniform( + "w", lower=0, upper=pt.square(z), transform=transform, default_transform=None + ) rvs = [x, y, z, w] if reversed: From 7d6ecf9a82309aa1a7524168806f9e0e3a2b3109 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 9 Apr 2024 18:32:55 +0300 Subject: [PATCH 09/24] add tests for transfor args --- tests/model/test_core.py | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index c5eaec430e..c4ac099df2 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -42,7 +42,7 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.distributions import Normal, transforms from pymc.distributions.distribution import PartialObservedRV -from pymc.distributions.transforms import log, simplex +from pymc.distributions.transforms import Transform, log, simplex from pymc.exceptions import ImputationWarning, ShapeError, ShapeWarning from pymc.logprob.basic import transformed_conditional_logp from pymc.logprob.transforms import IntervalTransform @@ -527,6 +527,42 @@ def test_model_var_maps(): assert model.rvs_to_transforms[x] is None +class TestTransformArgs: + def test_transform_warning(self): + with pm.Model(): + with pytest.warns( + UserWarning, + match="To disable default transform," + " please use default_transform=None" + " instead of transform=None. Setting transform to" + " None will not have any effect in future.", + ): + a = pm.Normal("a", transform=None) + + def test_transform_order(self): + transform_order = [] + + class DummyTransform(Transform): + name = "dummy1" + ndim_supp = 0 + + def __init__(self, marker) -> None: + super().__init__() + self.marker = marker + + def forward(self, value, *inputs): + nonlocal transform_order + transform_order.append(self.marker) + return value + + def backward(self, value, *inputs): + return value + + with pm.Model() as model: + x = pm.Normal("x", transform=DummyTransform(2), default_transform=DummyTransform(1)) + assert transform_order == [1, 2] + + def test_make_obs_var(): """ Check returned values for `data` given known inputs to `as_tensor()`. From 704aac68d354a5b43e67d89943fd56260edf88f7 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 9 Apr 2024 19:13:20 +0300 Subject: [PATCH 10/24] fix formating --- pymc/model/fgraph.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pymc/model/fgraph.py b/pymc/model/fgraph.py index b71c168f08..7a3fdd9829 100644 --- a/pymc/model/fgraph.py +++ b/pymc/model/fgraph.py @@ -320,7 +320,9 @@ def first_non_model_var(var): var, value, *dims = model_var.owner.inputs transform = model_var.owner.op.transform model.free_RVs.append(var) - model.create_value_var(var, transform=transform, default_transform=None, value_var=value) + model.create_value_var( + var, transform=transform, default_transform=None, value_var=value + ) model.set_initval(var, initval=None) elif isinstance(model_var.owner.op, ModelObservedRV): var, value, *dims = model_var.owner.inputs From 23bd69ffc974e7f44d9720bffffd9f62882c33b2 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Mon, 15 Apr 2024 20:17:20 +0400 Subject: [PATCH 11/24] Update pymc/model/core.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- pymc/model/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index db019b8008..1c9731efea 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1430,7 +1430,7 @@ def create_value_var( if transform is UNSET: transform = default_transform - elif transform and default_transform: + elif transform is not None and default_transform is not None: transform = ChainedTransform([default_transform, transform]) if value_var is None: From 35b44fed19b953acc2b1d39ae87dcbc8518e9845 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Mon, 15 Apr 2024 20:18:03 +0400 Subject: [PATCH 12/24] Update tests/distributions/test_transform.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/distributions/test_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_transform.py b/tests/distributions/test_transform.py index 4bc5e8543d..2b0f145bbe 100644 --- a/tests/distributions/test_transform.py +++ b/tests/distributions/test_transform.py @@ -619,7 +619,7 @@ def test_transform_univariate_dist_logp_shape(): def test_univariate_transform_multivariate_dist_raises(): with pm.Model() as m: - pm.Dirichlet("x", [1, 1, 1], transform=tr.log, default_transform=None) + pm.Dirichlet("x", [1, 1, 1], default_transform=tr.log) for jacobian_val in (True, False): with pytest.raises( From 7f0d1d3e97c7f48560be751bb492b0f9e17f5834 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Mon, 15 Apr 2024 20:18:31 +0400 Subject: [PATCH 13/24] Update tests/distributions/test_transform.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/distributions/test_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_transform.py b/tests/distributions/test_transform.py index 2b0f145bbe..8d464f206a 100644 --- a/tests/distributions/test_transform.py +++ b/tests/distributions/test_transform.py @@ -645,7 +645,7 @@ def log_jac_det(self, value, *inputs): buggy_transform = BuggyTransform() with pm.Model() as m: - pm.Uniform("x", shape=(4, 3), transform=buggy_transform, default_transform=None) + pm.Uniform("x", shape=(4, 3), default_transform=buggy_transform) for jacobian_val in (True, False): with pytest.raises( From 80641f89a0fcee84793124e8b3bdd01651f8e96e Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 16 Apr 2024 10:55:40 +0400 Subject: [PATCH 14/24] pass transform to default transform --- tests/logprob/test_utils.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/logprob/test_utils.py b/tests/logprob/test_utils.py index 7d3094c5b5..47cb65f195 100644 --- a/tests/logprob/test_utils.py +++ b/tests/logprob/test_utils.py @@ -218,15 +218,11 @@ def test_interdependent_transformed_rvs(self, reversed): transform = pm.distributions.transforms.Interval( bounds_fn=lambda *inputs: (inputs[-2], inputs[-1]) ) - x = pm.Uniform("x", lower=0, upper=1, transform=transform, default_transform=None) + x = pm.Uniform("x", lower=0, upper=1, default_transform=transform) # Operation between the variables provides a regression test for #7054 - y = pm.Uniform( - "y", lower=0, upper=pt.exp(x), transform=transform, default_transform=None - ) - z = pm.Uniform("z", lower=0, upper=y, transform=transform, default_transform=None) - w = pm.Uniform( - "w", lower=0, upper=pt.square(z), transform=transform, default_transform=None - ) + y = pm.Uniform("y", lower=0, upper=pt.exp(x), default_transform=transform) + z = pm.Uniform("z", lower=0, upper=y, default_transform=transform) + w = pm.Uniform("w", lower=0, upper=pt.square(z), default_transform=transform) rvs = [x, y, z, w] if reversed: From 80fa510db9d2e001e6b84cbf1a2f144e317e5c68 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 16 Apr 2024 10:59:09 +0400 Subject: [PATCH 15/24] remove duplicated test case --- tests/distributions/test_mixture.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index 97122ea6f8..e328c96931 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -1361,14 +1361,6 @@ def test_warning(self): warnings.simplefilter("error") Mixture("mix4", w=[0.5, 0.5], comp_dists=comp_dists, default_transform=None) - with pytest.warns( - UserWarning, - match="To disable default transform, please use default_transform=None" - " instead of transform=None. Setting transform to None will not have" - " any effect in future.", - ): - Mixture("mix5", w=[0.5, 0.5], comp_dists=comp_dists, transform=None) - with warnings.catch_warnings(): warnings.simplefilter("error") Mixture("mix6", w=[0.5, 0.5], comp_dists=comp_dists, observed=1) From 655a364edde066c6e13f26b1221095e8066b5f2f Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Tue, 16 Apr 2024 11:46:04 +0400 Subject: [PATCH 16/24] move transform=None warning to core module --- pymc/distributions/distribution.py | 9 --------- pymc/model/core.py | 9 +++++++++ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 36711383dd..3af0bb110e 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -398,15 +398,6 @@ def __new__( if not isinstance(name, string_types): raise TypeError(f"Name needs to be a string but got: {name}") - if transform is None and default_transform is UNSET: - default_transform = None - warnings.warn( - "To disable default transform, please use default_transform=None" - " instead of transform=None. Setting transform to None will" - " not have any effect in future.", - UserWarning, - ) - dims = convert_dims(dims) if observed is not None: observed = convert_observed_data(observed) diff --git a/pymc/model/core.py b/pymc/model/core.py index 1c9731efea..345e79d24a 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1422,6 +1422,15 @@ def create_value_var( # Make the value variable a transformed value variable, # if there's an applicable transform + if transform is None and default_transform is UNSET: + default_transform = None + warnings.warn( + "To disable default transform, please use default_transform=None" + " instead of transform=None. Setting transform to None will" + " not have any effect in future.", + UserWarning, + ) + if default_transform is UNSET: if rv_var.owner is None: default_transform = None From 995ec9337209edb0f279bfe3289408be31af94d4 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Wed, 17 Apr 2024 18:50:16 +0400 Subject: [PATCH 17/24] update type hints and docstrings --- pymc/model/core.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 345e79d24a..5abc5c65d8 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -22,7 +22,6 @@ from sys import modules from typing import ( TYPE_CHECKING, - Any, Literal, Optional, TypeVar, @@ -58,6 +57,7 @@ ) from pymc.initial_point import make_initial_point_fn from pymc.logprob.basic import transformed_conditional_logp +from pymc.logprob.transforms import Transform from pymc.logprob.utils import ParameterValueError, replace_rvs_by_values from pymc.model_graph import model_to_graphviz from pymc.pytensorf import ( @@ -1239,7 +1239,9 @@ def register_rv( dims : tuple Dimension names for the variable. transform - A transform for the random variable in log-likelihood space. + Additianal transform which may be applied after default transform. + default_transform + A default transform for the random variable in log-likelihood space. initval The initial value of the random variable. @@ -1298,8 +1300,8 @@ def make_obs_var( rv_var: TensorVariable, data: np.ndarray, dims, - transform: Any | None, - default_transform: Any | None, + transform: Transform | None, + default_transform: Transform | None, total_size: int | None, ) -> TensorVariable: """Create a `TensorVariable` for an observed random variable. @@ -1313,7 +1315,9 @@ def make_obs_var( The observed data. dims : tuple Dimension names for the variable. - transform : int, optional + transform + Additianal transform which may be applied after default transform. + default_transform A transform for the random variable in log-likelihood space. Returns @@ -1394,8 +1398,8 @@ def create_value_var( self, rv_var: TensorVariable, *, - transform: Any, - default_transform: Any, + transform: Transform, + default_transform: Transform, value_var: Variable | None = None, ) -> TensorVariable: """Create a ``TensorVariable`` that will be used as the random @@ -1411,7 +1415,11 @@ def create_value_var( ---------- rv_var : TensorVariable - transform : Any + transform: Transform + Additianal transform which may be applied after default transform. + + default_transform: Transform + A transform for the random variable in log-likelihood space. value_var : Variable, optional From 3eefa8e4a3d7dc0bc63c60fd3db7d8194837e353 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Wed, 17 Apr 2024 20:26:00 +0400 Subject: [PATCH 18/24] change test for transforms order --- tests/model/test_core.py | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index c4ac099df2..5df3b59460 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -42,7 +42,7 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.distributions import Normal, transforms from pymc.distributions.distribution import PartialObservedRV -from pymc.distributions.transforms import Transform, log, simplex +from pymc.distributions.transforms import Interval, LogTransform, log, simplex from pymc.exceptions import ImputationWarning, ShapeError, ShapeWarning from pymc.logprob.basic import transformed_conditional_logp from pymc.logprob.transforms import IntervalTransform @@ -540,27 +540,10 @@ def test_transform_warning(self): a = pm.Normal("a", transform=None) def test_transform_order(self): - transform_order = [] - - class DummyTransform(Transform): - name = "dummy1" - ndim_supp = 0 - - def __init__(self, marker) -> None: - super().__init__() - self.marker = marker - - def forward(self, value, *inputs): - nonlocal transform_order - transform_order.append(self.marker) - return value - - def backward(self, value, *inputs): - return value - with pm.Model() as model: - x = pm.Normal("x", transform=DummyTransform(2), default_transform=DummyTransform(1)) - assert transform_order == [1, 2] + x = pm.Normal("x", transform=Interval(0, 1), default_transform=log) + assert isinstance(model.rvs_to_transforms[x].transform_list[0], LogTransform) + assert isinstance(model.rvs_to_transforms[x].transform_list[1], Interval) def test_make_obs_var(): From 8a005fea2a8f0699284a72bfbe3270bb098473b4 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Thu, 18 Apr 2024 13:50:17 +0400 Subject: [PATCH 19/24] add new test case --- tests/model/test_core.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 5df3b59460..95383bf79a 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -545,6 +545,14 @@ def test_transform_order(self): assert isinstance(model.rvs_to_transforms[x].transform_list[0], LogTransform) assert isinstance(model.rvs_to_transforms[x].transform_list[1], Interval) + def test_default_transform_is_applied(self): + with pm.Model() as model1: + x1 = pm.LogNormal("x1", 0, 1, transform=Interval(-2, 2), default_transform=None) + with pm.Model() as model3: + x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log) + assert np.isinf(model1.compile_logp()({"x1_interval__": -1})) + assert np.isfinite(model3.compile_logp()({"x2_chain__": -1})) + def test_make_obs_var(): """ From 8a708f74615df316e1fbdd320cc766b477611a95 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Thu, 18 Apr 2024 16:13:33 +0400 Subject: [PATCH 20/24] Update pymc/model/core.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- pymc/model/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 5abc5c65d8..4e39f7b049 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1239,7 +1239,7 @@ def register_rv( dims : tuple Dimension names for the variable. transform - Additianal transform which may be applied after default transform. + Additional transform which may be applied after default transform. default_transform A default transform for the random variable in log-likelihood space. initval From 412ef7fc166f776b1d5d47fa7fb93a0fe0ce7f9f Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Thu, 18 Apr 2024 19:17:11 +0400 Subject: [PATCH 21/24] change test casses for transform args --- tests/model/test_core.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 95383bf79a..9bc436ef1b 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -42,7 +42,14 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.distributions import Normal, transforms from pymc.distributions.distribution import PartialObservedRV -from pymc.distributions.transforms import Interval, LogTransform, log, simplex +from pymc.distributions.transforms import ( + ChainedTransform, + Interval, + LogTransform, + log, + ordered, + simplex, +) from pymc.exceptions import ImputationWarning, ShapeError, ShapeWarning from pymc.logprob.basic import transformed_conditional_logp from pymc.logprob.transforms import IntervalTransform @@ -542,16 +549,17 @@ def test_transform_warning(self): def test_transform_order(self): with pm.Model() as model: x = pm.Normal("x", transform=Interval(0, 1), default_transform=log) + assert isinstance(model.rvs_to_transforms[x], ChainedTransform) assert isinstance(model.rvs_to_transforms[x].transform_list[0], LogTransform) assert isinstance(model.rvs_to_transforms[x].transform_list[1], Interval) def test_default_transform_is_applied(self): with pm.Model() as model1: - x1 = pm.LogNormal("x1", 0, 1, transform=Interval(-2, 2), default_transform=None) + x1 = pm.LogNormal("x1", [0, 0], [1, 1], transform=ordered, default_transform=None) with pm.Model() as model3: - x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log) - assert np.isinf(model1.compile_logp()({"x1_interval__": -1})) - assert np.isfinite(model3.compile_logp()({"x2_chain__": -1})) + x2 = pm.LogNormal("x2", [0, 0], [1, 1], transform=ordered, default_transform=log) + assert np.isinf(model1.compile_logp()({"x1_ordered__": (-1, -1)})) + assert np.isfinite(model3.compile_logp()({"x2_chain__": (-1, -1)})) def test_make_obs_var(): From 48f5f8577c8b22e6188cb66ced4b76b91e86c721 Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Thu, 18 Apr 2024 19:24:56 +0400 Subject: [PATCH 22/24] fix docstring and change argument order --- pymc/model/core.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 4e39f7b049..2f73c9ee24 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -1221,8 +1221,8 @@ def register_rv( observed=None, total_size=None, dims=None, - transform=UNSET, default_transform=UNSET, + transform=UNSET, initval=None, ): """Register an (un)observed random variable with the model. @@ -1238,10 +1238,10 @@ def register_rv( upscales logp of variable with ``coef = total_size/var.shape[0]`` dims : tuple Dimension names for the variable. - transform - Additional transform which may be applied after default transform. default_transform A default transform for the random variable in log-likelihood space. + transform + Additional transform which may be applied after default transform. initval The initial value of the random variable. @@ -1290,7 +1290,7 @@ def register_rv( # `rv_var` is potentially changed by `make_obs_var`, # for example into a new graph for imputation of missing data. rv_var = self.make_obs_var( - rv_var, observed, dims, transform, default_transform, total_size + rv_var, observed, dims, default_transform, transform, total_size ) return rv_var @@ -1300,8 +1300,8 @@ def make_obs_var( rv_var: TensorVariable, data: np.ndarray, dims, - transform: Transform | None, default_transform: Transform | None, + transform: Transform | None, total_size: int | None, ) -> TensorVariable: """Create a `TensorVariable` for an observed random variable. @@ -1315,10 +1315,10 @@ def make_obs_var( The observed data. dims : tuple Dimension names for the variable. - transform - Additianal transform which may be applied after default transform. default_transform A transform for the random variable in log-likelihood space. + transform + Additional transform which may be applied after default transform. Returns ------- @@ -1398,8 +1398,8 @@ def create_value_var( self, rv_var: TensorVariable, *, - transform: Transform, default_transform: Transform, + transform: Transform, value_var: Variable | None = None, ) -> TensorVariable: """Create a ``TensorVariable`` that will be used as the random @@ -1415,12 +1415,12 @@ def create_value_var( ---------- rv_var : TensorVariable - transform: Transform - Additianal transform which may be applied after default transform. - default_transform: Transform A transform for the random variable in log-likelihood space. + transform: Transform + Additional transform which may be applied after default transform. + value_var : Variable, optional Returns From bff43718b9c61edec544171ac2fef0023bd36558 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Thu, 18 Apr 2024 20:24:09 +0400 Subject: [PATCH 23/24] Apply suggestions from code review Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/model/test_core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 9bc436ef1b..5d33e47e57 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -556,8 +556,8 @@ def test_transform_order(self): def test_default_transform_is_applied(self): with pm.Model() as model1: x1 = pm.LogNormal("x1", [0, 0], [1, 1], transform=ordered, default_transform=None) - with pm.Model() as model3: - x2 = pm.LogNormal("x2", [0, 0], [1, 1], transform=ordered, default_transform=log) + with pm.Model() as model2: + x2 = pm.LogNormal("x2", [0, 0], [1, 1], transform=ordered) assert np.isinf(model1.compile_logp()({"x1_ordered__": (-1, -1)})) assert np.isfinite(model3.compile_logp()({"x2_chain__": (-1, -1)})) From e0f26b099f1fb4af56efee491657e5b258363dcb Mon Sep 17 00:00:00 2001 From: Anatoly Rubanov Date: Thu, 18 Apr 2024 20:27:04 +0400 Subject: [PATCH 24/24] refactor test case --- tests/model/test_core.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/model/test_core.py b/tests/model/test_core.py index 5d33e47e57..bb68d07886 100644 --- a/tests/model/test_core.py +++ b/tests/model/test_core.py @@ -549,9 +549,10 @@ def test_transform_warning(self): def test_transform_order(self): with pm.Model() as model: x = pm.Normal("x", transform=Interval(0, 1), default_transform=log) - assert isinstance(model.rvs_to_transforms[x], ChainedTransform) - assert isinstance(model.rvs_to_transforms[x].transform_list[0], LogTransform) - assert isinstance(model.rvs_to_transforms[x].transform_list[1], Interval) + transform = model.rvs_to_transforms[x] + assert isinstance(transform, ChainedTransform) + assert isinstance(transform.transform_list[0], LogTransform) + assert isinstance(transform.transform_list[1], Interval) def test_default_transform_is_applied(self): with pm.Model() as model1: @@ -559,7 +560,7 @@ def test_default_transform_is_applied(self): with pm.Model() as model2: x2 = pm.LogNormal("x2", [0, 0], [1, 1], transform=ordered) assert np.isinf(model1.compile_logp()({"x1_ordered__": (-1, -1)})) - assert np.isfinite(model3.compile_logp()({"x2_chain__": (-1, -1)})) + assert np.isfinite(model2.compile_logp()({"x2_chain__": (-1, -1)})) def test_make_obs_var():