From e67a317af775af8c864e95e8806f40b51c1a39b6 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 10:36:37 +0100 Subject: [PATCH 01/11] change default transform LKJCOrr --- pymc/distributions/multivariate.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 42a2a9f62c..6cbd9011e5 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -1524,6 +1524,13 @@ def _random_corr_matrix(cls, rng, n, eta, flat_size): lkjcorr = LKJCorrRV() +class MultivariateIntervalTransform(Interval): + name = "interval" + + def log_jac_det(self, *args): + return super().log_jac_det(*args).sum(-1) + + class LKJCorr(BoundedContinuous): r""" The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood. @@ -1623,7 +1630,7 @@ def logp(value, n, eta): @_default_transform.register(LKJCorr) def lkjcorr_default_transform(op, rv): - return Interval(floatX(-1.0), floatX(1.0)) + return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0)) class MatrixNormalRV(RandomVariable): From 301b1591246f552b81a2b28fe0936420b42aa14e Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 10:48:19 +0100 Subject: [PATCH 02/11] add test --- tests/distributions/test_multivariate.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 027d0b1915..48967e17ee 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2121,6 +2121,12 @@ def ref_rand(size, n, eta): size=1000, ) + @pytest.mark.parametrize(argnames="n", argvalues=[2, 3], ids=["n=2", "n=3"]) + def test_default_transform(self, n): + with pm.Model() as m: + pm.LKJCorr("x", n=n, eta=1) + m.logp() + class TestLKJCholeskyCov(BaseTestDistributionRandom): pymc_dist = _LKJCholeskyCov From bc0d8f55ad15cd150a0b5ce74fcebe4839f5f5ab Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Wed, 13 Dec 2023 11:20:54 +0100 Subject: [PATCH 03/11] Update tests/distributions/test_multivariate.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/distributions/test_multivariate.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 48967e17ee..e8856dd1a5 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2121,11 +2121,11 @@ def ref_rand(size, n, eta): size=1000, ) - @pytest.mark.parametrize(argnames="n", argvalues=[2, 3], ids=["n=2", "n=3"]) - def test_default_transform(self, n): + def test_default_transform(self): with pm.Model() as m: - pm.LKJCorr("x", n=n, eta=1) - m.logp() + x = pm.LKJCorr("x", n=2, eta=1, shape=(3, 2)) + assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) + assert m.logp(sum=False)[0].shape == (3,) class TestLKJCholeskyCov(BaseTestDistributionRandom): From 9b4bf2acf82798d13cf3d68dbcade82d5dfca765 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 11:30:21 +0100 Subject: [PATCH 04/11] fix import --- tests/distributions/test_multivariate.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index e8856dd1a5..f2c4a0d4b4 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -33,6 +33,7 @@ import pymc as pm from pymc.distributions.multivariate import ( + MultivariateIntervalTransform, _LKJCholeskyCov, _OrderedMultinomial, posdef, From 21948a440797c3296408f2798c4db3b56a7cf742 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 12:22:08 +0100 Subject: [PATCH 05/11] error handling --- pymc/logprob/transform_value.py | 13 ++++++------- tests/distributions/test_multivariate.py | 21 +++++++++++++++++---- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/pymc/logprob/transform_value.py b/pymc/logprob/transform_value.py index ee37fdca1d..4bcf130f76 100644 --- a/pymc/logprob/transform_value.py +++ b/pymc/logprob/transform_value.py @@ -125,13 +125,12 @@ def transformed_value_logprob(op, values, *rv_outs, use_jacobian=True, **kwargs) raise NotImplementedError( f"Univariate transform {transform} cannot be applied to multivariate {rv_op}" ) - else: - # Check there is no broadcasting between logp and jacobian - if logp.type.broadcastable != log_jac_det.type.broadcastable: - raise ValueError( - f"The logp of {rv_op} and log_jac_det of {transform} are not allowed to broadcast together. " - "There is a bug in the implementation of either one." - ) + # Check there is no broadcasting between logp and jacobian + if logp.type.broadcastable != log_jac_det.type.broadcastable: + raise ValueError( + f"The logp of {rv_op} and log_jac_det of {transform} are not allowed to broadcast together. " + "There is a bug in the implementation of either one." + ) if use_jacobian: if value.name: diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index f2c4a0d4b4..64d8df3917 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -13,7 +13,6 @@ # limitations under the License. import functools as ft -import re import warnings import numpy as np @@ -2122,11 +2121,25 @@ def ref_rand(size, n, eta): size=1000, ) - def test_default_transform(self): + @pytest.mark.parametrize( + argnames="shape, expected_shape", + argvalues=[ + ((2,), ()), + pytest.param( + (3, 2), + (3,), + marks=pytest.mark.xfail( + raises=NotImplementedError, + reason="We do not support batch dimensions for pm.LKJCorr yet.", + ), + ), + ], + ) + def test_default_transform(self, shape, expected_shape): with pm.Model() as m: - x = pm.LKJCorr("x", n=2, eta=1, shape=(3, 2)) + x = pm.LKJCorr("x", n=2, eta=1, shape=shape) assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) - assert m.logp(sum=False)[0].shape == (3,) + assert m.logp(sum=False)[0].type.shape == expected_shape class TestLKJCholeskyCov(BaseTestDistributionRandom): From 8e76b8457b51b286acd191765d6bade7bad639b8 Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Wed, 13 Dec 2023 12:48:52 +0100 Subject: [PATCH 06/11] Update tests/distributions/test_multivariate.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/distributions/test_multivariate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 64d8df3917..933da67c89 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2139,7 +2139,7 @@ def test_default_transform(self, shape, expected_shape): with pm.Model() as m: x = pm.LKJCorr("x", n=2, eta=1, shape=shape) assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) - assert m.logp(sum=False)[0].type.shape == expected_shape + assert m.logp(sum=False)[0].type.shape == shape[:-1] class TestLKJCholeskyCov(BaseTestDistributionRandom): From 6f03ec55ed366abef6b48006dfb1de49023d5e56 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 12:51:15 +0100 Subject: [PATCH 07/11] remove expected shapre argument --- tests/distributions/test_multivariate.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 933da67c89..aec678feee 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2122,12 +2122,11 @@ def ref_rand(size, n, eta): ) @pytest.mark.parametrize( - argnames="shape, expected_shape", + argnames="shape", argvalues=[ - ((2,), ()), + (2,), pytest.param( (3, 2), - (3,), marks=pytest.mark.xfail( raises=NotImplementedError, reason="We do not support batch dimensions for pm.LKJCorr yet.", @@ -2135,7 +2134,7 @@ def ref_rand(size, n, eta): ), ], ) - def test_default_transform(self, shape, expected_shape): + def test_default_transform(self, shape): with pm.Model() as m: x = pm.LKJCorr("x", n=2, eta=1, shape=shape) assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) From aa7810c52b6a761c40a60c9ed17282fcf00f8d14 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 12:56:57 +0100 Subject: [PATCH 08/11] add notimplemented error for logp --- pymc/distributions/multivariate.py | 3 +++ tests/distributions/test_multivariate.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 6cbd9011e5..ca1760c387 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -1599,6 +1599,9 @@ def logp(value, n, eta): TensorVariable """ + if value.ndim > 1: + raise NotImplementedError("LKJCorr logp is only implemented for vector values (ndim=1)") + # TODO: PyTensor does not have a `triu_indices`, so we can only work with constant # n (or else find a different expression) if not isinstance(n, Constant): diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index aec678feee..dc237d4211 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2129,7 +2129,7 @@ def ref_rand(size, n, eta): (3, 2), marks=pytest.mark.xfail( raises=NotImplementedError, - reason="We do not support batch dimensions for pm.LKJCorr yet.", + reason="LKJCorr logp is only implemented for vector values (ndim=1)", ), ), ], From bee4144547b67b584c18c9ebfb14e9a6ea64fc83 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 13:03:21 +0100 Subject: [PATCH 09/11] make standalone test --- tests/distributions/test_multivariate.py | 35 ++++++++++++------------ 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index dc237d4211..9fd40dcb27 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2121,24 +2121,25 @@ def ref_rand(size, n, eta): size=1000, ) - @pytest.mark.parametrize( - argnames="shape", - argvalues=[ - (2,), - pytest.param( - (3, 2), - marks=pytest.mark.xfail( - raises=NotImplementedError, - reason="LKJCorr logp is only implemented for vector values (ndim=1)", - ), + +@pytest.mark.parametrize( + argnames="shape", + argvalues=[ + (2,), + pytest.param( + (3, 2), + marks=pytest.mark.xfail( + raises=NotImplementedError, + reason="LKJCorr logp is only implemented for vector values (ndim=1)", ), - ], - ) - def test_default_transform(self, shape): - with pm.Model() as m: - x = pm.LKJCorr("x", n=2, eta=1, shape=shape) - assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) - assert m.logp(sum=False)[0].type.shape == shape[:-1] + ), + ], +) +def test_default_transform(shape): + with pm.Model() as m: + x = pm.LKJCorr("x", n=2, eta=1, shape=shape) + assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) + assert m.logp(sum=False)[0].type.shape == shape[:-1] class TestLKJCholeskyCov(BaseTestDistributionRandom): From 67954c776ee54d064b0d2c2e566ef9148892665b Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Wed, 13 Dec 2023 13:07:27 +0100 Subject: [PATCH 10/11] Update tests/distributions/test_multivariate.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/distributions/test_multivariate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 9fd40dcb27..f546747013 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -2135,7 +2135,7 @@ def ref_rand(size, n, eta): ), ], ) -def test_default_transform(shape): +def test_LKJCorr_default_transform(shape): with pm.Model() as m: x = pm.LKJCorr("x", n=2, eta=1, shape=shape) assert isinstance(m.rvs_to_transforms[x], MultivariateIntervalTransform) From 67b80c66deaf692b529b18b54a6da609dfb63896 Mon Sep 17 00:00:00 2001 From: juanitorduz Date: Wed, 13 Dec 2023 14:29:44 +0100 Subject: [PATCH 11/11] patch test --- tests/distributions/test_multivariate.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index f546747013..2390804c0f 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -1306,8 +1306,26 @@ def test_kronecker_normal_moment(self, mu, covs, size, expected): [ (3, 1, None, np.zeros(3)), (5, 1, None, np.zeros(10)), - (3, 1, 1, np.zeros((1, 3))), - (5, 1, (2, 3), np.zeros((2, 3, 10))), + pytest.param( + 3, + 1, + 1, + np.zeros((1, 3)), + marks=pytest.mark.xfail( + raises=NotImplementedError, + reason="LKJCorr logp is only implemented for vector values (ndim=1)", + ), + ), + pytest.param( + 5, + 1, + (2, 3), + np.zeros((2, 3, 10)), + marks=pytest.mark.xfail( + raises=NotImplementedError, + reason="LKJCorr logp is only implemented for vector values (ndim=1)", + ), + ), ], ) def test_lkjcorr_moment(self, n, eta, size, expected):