From b4546220b1f7fb971b55379b3a055bb1dfae5b05 Mon Sep 17 00:00:00 2001 From: Ricardo Date: Tue, 19 Apr 2022 17:28:17 +0200 Subject: [PATCH 1/5] fix test_step_args wrongly being inside the body of a class --- pymc/tests/test_sampling.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/tests/test_sampling.py b/pymc/tests/test_sampling.py index 319b656a9e..8e3550021e 100644 --- a/pymc/tests/test_sampling.py +++ b/pymc/tests/test_sampling.py @@ -1406,8 +1406,8 @@ def test_draw_aesara_function_kwargs(self): assert np.all(draws == np.arange(5)) -class test_step_args(SeededTest): - with pm.Model() as model: +def test_step_args(): + with pm.Model(rng_seeder=1410) as model: a = pm.Normal("a") idata0 = pm.sample(target_accept=0.5) idata1 = pm.sample(nuts={"target_accept": 0.5}) @@ -1415,7 +1415,7 @@ class test_step_args(SeededTest): npt.assert_almost_equal(idata0.sample_stats.acceptance_rate.mean(), 0.5, decimal=1) npt.assert_almost_equal(idata1.sample_stats.acceptance_rate.mean(), 0.5, decimal=1) - with pm.Model() as model: + with pm.Model(rng_seeder=1418) as model: a = pm.Normal("a") b = pm.Poisson("b", 1) idata0 = pm.sample(target_accept=0.5) From e998498c344dfac8bb84809c5e20860fb9f7118b Mon Sep 17 00:00:00 2001 From: Ricardo Date: Tue, 19 Apr 2022 11:55:41 +0200 Subject: [PATCH 2/5] Remove unused variable in GaussianRandomWalk.rng_fn --- pymc/distributions/timeseries.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index bf6addb900..60dc4d8103 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -116,7 +116,6 @@ def rng_fn( # If size is None then the returned series should be (*size, 1+steps) else: - init_size = (*size, 1) dist_shape = (*size, int(steps)) innovations = rng.normal(loc=mu, scale=sigma, size=dist_shape) From a9224d6e7d92800217011177c2ea5f64c8b1f348 Mon Sep 17 00:00:00 2001 From: Ricardo Date: Tue, 19 Apr 2022 09:36:12 +0200 Subject: [PATCH 3/5] Broadcast dist in Censored with lower and upper when size is not specified --- pymc/distributions/censored.py | 20 ++++++++++---------- pymc/tests/test_distributions.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/pymc/distributions/censored.py b/pymc/distributions/censored.py index 60edd7e227..c0790dee3e 100644 --- a/pymc/distributions/censored.py +++ b/pymc/distributions/censored.py @@ -18,6 +18,7 @@ from aesara.tensor import TensorVariable from aesara.tensor.random.op import RandomVariable +from pymc.aesaraf import change_rv_size from pymc.distributions.distribution import SymbolicDistribution, _moment from pymc.util import check_dist_not_registered @@ -74,10 +75,13 @@ def dist(cls, dist, lower, upper, **kwargs): @classmethod def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None): - if lower is None: - lower = at.constant(-np.inf) - if upper is None: - upper = at.constant(np.inf) + + lower = at.constant(-np.inf) if lower is None else at.as_tensor_variable(lower) + upper = at.constant(np.inf) if upper is None else at.as_tensor_variable(upper) + + # When size is not specified, dist may have to be broadcasted according to lower/upper + dist_shape = size if size is not None else at.broadcast_shape(dist, lower, upper) + dist = change_rv_size(dist, dist_shape) # Censoring is achieved by clipping the base distribution between lower and upper rv_out = at.clip(dist, lower, upper) @@ -88,8 +92,6 @@ def rv_op(cls, dist, lower=None, upper=None, size=None, rngs=None): rv_out.tag.lower = lower rv_out.tag.upper = upper - if size is not None: - rv_out = cls.change_size(rv_out, size) if rngs is not None: rv_out = cls.change_rngs(rv_out, rngs) @@ -101,12 +103,10 @@ def ndim_supp(cls, *dist_params): @classmethod def change_size(cls, rv, new_size, expand=False): - dist_node = rv.tag.dist.owner + dist = rv.tag.dist lower = rv.tag.lower upper = rv.tag.upper - rng, old_size, dtype, *dist_params = dist_node.inputs - new_size = new_size if not expand else tuple(new_size) + tuple(old_size) - new_dist = dist_node.op.make_node(rng, new_size, dtype, *dist_params).default_output() + new_dist = change_rv_size(dist, new_size, expand=expand) return cls.rv_op(new_dist, lower, upper) @classmethod diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index 661b05d91e..20a2bdaa3d 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -3381,6 +3381,18 @@ def test_change_size(self): new_dist = pm.Censored.change_size(base_dist, (4,), expand=True) assert new_dist.eval().shape == (4, 3, 2) + def test_dist_broadcasted_by_lower_upper(self): + x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=None) + assert tuple(x.owner.inputs[0].shape.eval()) == (2,) + + x = pm.Censored.dist(pm.Normal.dist(), lower=np.zeros((2,)), upper=np.zeros((4, 2))) + assert tuple(x.owner.inputs[0].shape.eval()) == (4, 2) + + x = pm.Censored.dist( + pm.Normal.dist(size=(3, 4, 2)), lower=np.zeros((2,)), upper=np.zeros((4, 2)) + ) + assert tuple(x.owner.inputs[0].shape.eval()) == (3, 4, 2) + class TestLKJCholeskCov: def test_dist(self): From 95cc466ed79f6aa68a9e0429dea5d881f4caf9fe Mon Sep 17 00:00:00 2001 From: Ricardo Date: Tue, 19 Apr 2022 10:24:54 +0200 Subject: [PATCH 4/5] Move _LKJCholeskyCov sd_dist resizing to Op.make_node --- pymc/distributions/multivariate.py | 31 +++++++++++++++--------------- pymc/tests/test_distributions.py | 14 ++++++++++++-- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 354061bab7..940b3907e1 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -32,7 +32,7 @@ from aesara.tensor.nlinalg import det, eigh, matrix_inverse, trace from aesara.tensor.random.basic import MultinomialRV, dirichlet, multivariate_normal from aesara.tensor.random.op import RandomVariable, default_supp_shape_from_params -from aesara.tensor.random.utils import broadcast_params +from aesara.tensor.random.utils import broadcast_params, normalize_size_param from aesara.tensor.slinalg import Cholesky from aesara.tensor.slinalg import solve_lower_triangular as solve_lower from aesara.tensor.slinalg import solve_upper_triangular as solve_upper @@ -1134,6 +1134,19 @@ def make_node(self, rng, size, dtype, n, eta, D): D = at.as_tensor_variable(D) + # We resize the sd_dist `D` automatically so that it has (size x n) independent + # draws which is what the `_LKJCholeskyCovRV.rng_fn` expects. This makes the + # random and logp methods equivalent, as the latter also assumes a unique value + # for each diagonal element. + # Since `eta` and `n` are forced to be scalars we don't need to worry about + # implied batched dimensions for the time being. + size = normalize_size_param(size) + if D.owner.op.ndim_supp == 0: + D = change_rv_size(D, at.concatenate((size, (n,)))) + else: + # The support shape must be `n` but we have no way of controlling it + D = change_rv_size(D, size) + return super().make_node(rng, size, dtype, n, eta, D) def _infer_shape(self, size, dist_params, param_shapes=None): @@ -1179,7 +1192,7 @@ def __new__(cls, name, eta, n, sd_dist, **kwargs): return super().__new__(cls, name, eta, n, sd_dist, **kwargs) @classmethod - def dist(cls, eta, n, sd_dist, size=None, **kwargs): + def dist(cls, eta, n, sd_dist, **kwargs): eta = at.as_tensor_variable(floatX(eta)) n = at.as_tensor_variable(intX(n)) @@ -1191,18 +1204,6 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs): ): raise TypeError("sd_dist must be a scalar or vector distribution variable") - # We resize the sd_dist automatically so that it has (size x n) independent draws - # which is what the `_LKJCholeskyCovRV.rng_fn` expects. This makes the random - # and logp methods equivalent, as the latter also assumes a unique value for each - # diagonal element. - # Since `eta` and `n` are forced to be scalars we don't need to worry about - # implied batched dimensions for the time being. - if sd_dist.owner.op.ndim_supp == 0: - sd_dist = change_rv_size(sd_dist, to_tuple(size) + (n,)) - else: - # The support shape must be `n` but we have no way of controlling it - sd_dist = change_rv_size(sd_dist, to_tuple(size)) - # sd_dist is part of the generative graph, but should be completely ignored # by the logp graph, since the LKJ logp explicitly includes these terms. # Setting sd_dist.tag.ignore_logprob to True, will prevent Aeppl warning about @@ -1211,7 +1212,7 @@ def dist(cls, eta, n, sd_dist, size=None, **kwargs): # sd_dist prior components from the logp expression. sd_dist.tag.ignore_logprob = True - return super().dist([n, eta, sd_dist], size=size, **kwargs) + return super().dist([n, eta, sd_dist], **kwargs) def moment(rv, size, n, eta, sd_dists): diag_idxs = (at.cumsum(at.arange(1, n + 1)) - 1).astype("int32") diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index 20a2bdaa3d..cb485a086f 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -3437,8 +3437,18 @@ def test_no_warning_logp(self): pm.MvNormal.dist(np.ones(3), np.eye(3)), ], ) - def test_sd_dist_automatically_resized(self, sd_dist): - x = pm.LKJCholeskyCov.dist(n=3, eta=1, sd_dist=sd_dist, size=10, compute_corr=False) + @pytest.mark.parametrize( + "size, shape", + [ + ((10,), None), + (None, (10, 6)), + (None, (10, ...)), + ], + ) + def test_sd_dist_automatically_resized(self, sd_dist, size, shape): + x = pm.LKJCholeskyCov.dist( + n=3, eta=1, sd_dist=sd_dist, size=size, shape=shape, compute_corr=False + ) resized_sd_dist = x.owner.inputs[-1] assert resized_sd_dist.eval().shape == (10, 3) # LKJCov has support shape `(n * (n+1)) // 2` From fe4ec8ea24b6964ee2123fb4eddfc2eb1c3a4581 Mon Sep 17 00:00:00 2001 From: Ricardo Date: Tue, 19 Apr 2022 11:55:48 +0200 Subject: [PATCH 5/5] Move GaussianRandomWalk init resizing to Op.make_node --- pymc/distributions/timeseries.py | 28 ++++++++++----------- pymc/tests/test_distributions_timeseries.py | 6 +++++ 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 60dc4d8103..c7225ad50e 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -19,12 +19,13 @@ from aesara import scan from aesara.tensor.random.op import RandomVariable +from aesara.tensor.random.utils import normalize_size_param from pymc.aesaraf import change_rv_size, floatX, intX from pymc.distributions import distribution, logprob, multivariate from pymc.distributions.continuous import Flat, Normal, get_tau_sigma from pymc.distributions.dist_math import check_parameters -from pymc.distributions.shape_utils import to_tuple +from pymc.distributions.shape_utils import rv_size_is_none, to_tuple from pymc.util import check_dist_not_registered __all__ = [ @@ -54,6 +55,16 @@ def make_node(self, rng, size, dtype, mu, sigma, init, steps): if not steps.ndim == 0 or not steps.dtype.startswith("int"): raise ValueError("steps must be an integer scalar (ndim=0).") + mu = at.as_tensor_variable(mu) + sigma = at.as_tensor_variable(sigma) + init = at.as_tensor_variable(init) + + # Resize init distribution + size = normalize_size_param(size) + # If not explicit, size is determined by the shapes of mu, sigma, and init + init_size = size if not rv_size_is_none(size) else at.broadcast_shape(mu, sigma, init) + init = change_rv_size(init, init_size) + return super().make_node(rng, size, dtype, mu, sigma, init, steps) def _supp_shape_from_params(self, dist_params, reop_param_idx=0, param_shapes=None): @@ -160,15 +171,9 @@ def dist( raise ValueError("Must specify steps parameter") steps = at.as_tensor_variable(intX(steps)) - shape = kwargs.get("shape", None) - if size is None and shape is None: - init_size = None - else: - init_size = to_tuple(size) if size is not None else to_tuple(shape)[:-1] - # If no scalar distribution is passed then initialize with a Normal of same mu and sigma if init is None: - init = Normal.dist(mu, sigma, size=init_size) + init = Normal.dist(mu, sigma) else: if not ( isinstance(init, at.TensorVariable) @@ -178,13 +183,6 @@ def dist( ): raise TypeError("init must be a univariate distribution variable") - if init_size is not None: - init = change_rv_size(init, init_size) - else: - # If not explicit, size is determined by the shapes of mu, sigma, and init - bcast_shape = at.broadcast_arrays(mu, sigma, init)[0].shape - init = change_rv_size(init, bcast_shape) - # Ignores logprob of init var because that's accounted for in the logp method init.tag.ignore_logprob = True diff --git a/pymc/tests/test_distributions_timeseries.py b/pymc/tests/test_distributions_timeseries.py index 76bb91daeb..b0c5325bce 100644 --- a/pymc/tests/test_distributions_timeseries.py +++ b/pymc/tests/test_distributions_timeseries.py @@ -103,6 +103,12 @@ def test_gaussian_random_walk_init_dist_shape(init): assert tuple(grw.owner.inputs[-2].shape.eval()) == (3, 2) +def test_shape_ellipsis(): + grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=5, init=pm.Normal.dist(), shape=(3, ...)) + assert tuple(grw.shape.eval()) == (3, 6) + assert tuple(grw.owner.inputs[-2].shape.eval()) == (3,) + + def test_gaussianrandomwalk_broadcasted_by_init_dist(): grw = pm.GaussianRandomWalk.dist(mu=0, sigma=1, steps=4, init=pm.Normal.dist(size=(2, 3))) assert tuple(grw.shape.eval()) == (2, 3, 5)