From c6b920d85939ab24a4f5771e6f323dda3fd2c3fe Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sun, 25 Jan 2026 15:34:27 +0100 Subject: [PATCH 1/4] distribution module: remove eager dependency on model module --- pymc/distributions/custom.py | 7 +++++-- pymc/distributions/distribution.py | 7 ++++--- pymc/distributions/shape_utils.py | 3 ++- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/pymc/distributions/custom.py b/pymc/distributions/custom.py index 9def0aba2f..ce557bbc56 100644 --- a/pymc/distributions/custom.py +++ b/pymc/distributions/custom.py @@ -40,7 +40,6 @@ from pymc.distributions.shape_utils import _change_dist_size, rv_size_is_none from pymc.exceptions import BlockModelAccessError from pymc.logprob.abstract import _logcdf, _logprob -from pymc.model.core import new_or_existing_block_model_access from pymc.pytensorf import collect_default_updates @@ -265,6 +264,8 @@ def rv_op( class_name: str, rng=None, ): + from pymc.model.core import new_or_existing_block_model_access + size = normalize_size_param(size) # If it's NoneConst, just use that as the dummy dummy_size_param = size.type() if isinstance(size, TensorVariable) else size @@ -827,11 +828,13 @@ def check_valid_dist_random(cls, dist, random, dist_params): @classmethod def is_symbolic_random(self, random, dist_params): + from pymc.model.core import new_or_existing_block_model_access + if random is None: return False # Try calling random with symbolic inputs + size = normalize_size_param(None) try: - size = normalize_size_param(None) with new_or_existing_block_model_access( error_msg_on_access="Model variables cannot be created in the random function. Use the `.dist` API to create such variables." ): diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index dbf7625156..5ca12a2979 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -54,7 +54,6 @@ from pymc.logprob.abstract import MeasurableOp, _icdf, _logccdf, _logcdf, _logprob from pymc.logprob.basic import logp from pymc.logprob.rewriting import logprob_rewrites_db -from pymc.printing import str_for_dist from pymc.pytensorf import ( collect_default_updates_inner_fgraph, constant_fold, @@ -524,9 +523,9 @@ def __new__( rv : TensorVariable The created random variable tensor, registered in the Model. """ - try: - from pymc.model import Model + from pymc.model.core import Model + try: model = Model.get_context() except TypeError: raise TypeError( @@ -565,6 +564,8 @@ def __new__( ) # add in pretty-printing support + from pymc.printing import str_for_dist + rv_out.str_repr = types.MethodType(str_for_dist, rv_out) rv_out._repr_latex_ = types.MethodType( functools.partial(str_for_dist, formatting="latex"), rv_out diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 4b8131212b..c9319ab995 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -30,7 +30,6 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -from pymc.model import modelcontext from pymc.pytensorf import convert_observed_data, resolve_shapes __all__ = [ @@ -386,6 +385,8 @@ def get_support_shape( inferred_support_shape = [shape[i] - support_shape_offset[i] for i in range(-ndim_supp, 0)] if inferred_support_shape is None and dims is not None: + from pymc.model import modelcontext + dims = convert_dims(dims) assert isinstance(dims, tuple) if len(dims) < ndim_supp: From d15cb2dd4409b2bfc01e596bdb0f84db3dc92d13 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sun, 25 Jan 2026 15:34:59 +0100 Subject: [PATCH 2/4] Use cleaner variadic_[add|mul] in rewrite --- pymc/logprob/transforms.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pymc/logprob/transforms.py b/pymc/logprob/transforms.py index 093e0802f4..cfadd1d988 100644 --- a/pymc/logprob/transforms.py +++ b/pymc/logprob/transforms.py @@ -103,6 +103,8 @@ sub, tanh, true_div, + variadic_add, + variadic_mul, ) from pytensor.tensor.variable import TensorVariable @@ -518,12 +520,12 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Varia transform_inputs = (measurable_input, power) transform = PowerTransform(power=power) elif isinstance(scalar_op, Add): - transform_inputs = (measurable_input, pt.add(*other_inputs)) + transform_inputs = (measurable_input, variadic_add(*other_inputs)) transform = LocTransform( transform_args_fn=lambda *inputs: inputs[-1], ) elif isinstance(scalar_op, Mul): - transform_inputs = (measurable_input, pt.mul(*other_inputs)) + transform_inputs = (measurable_input, variadic_mul(*other_inputs)) transform = ScaleTransform( transform_args_fn=lambda *inputs: inputs[-1], ) From 4175fdaaaf4e9a500f4f4a995d0c51912d5da844 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sun, 25 Jan 2026 15:46:22 +0100 Subject: [PATCH 3/4] Organize test_order and remove dependency on PyMC distributions --- tests/logprob/test_order.py | 474 ++++++++++++++++++------------------ 1 file changed, 231 insertions(+), 243 deletions(-) diff --git a/tests/logprob/test_order.py b/tests/logprob/test_order.py index 32c5d3c001..d7300b18fc 100644 --- a/tests/logprob/test_order.py +++ b/tests/logprob/test_order.py @@ -42,253 +42,241 @@ import pytest import scipy.stats as sp -import pymc as pm - -from pymc import logp -from pymc.logprob import conditional_logp +from pymc.logprob import conditional_logp, logp from pymc.testing import assert_no_rvs -def test_argmax(): - """Test whether the logprob for ```pt.argmax``` is correctly rejected""" - x = pt.random.normal(0, 1, size=(3,)) - x.name = "x" - x_argmax = pt.argmax(x, axis=-1) - x_max_value = pt.scalar("x_max_value", dtype=x_argmax.type.dtype) - - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented for Argmax")): - logp(x_argmax, x_max_value) - - -@pytest.mark.parametrize( - "pt_op", - [ - pt.max, - pt.min, - ], -) -def test_non_iid_fails(pt_op): - """Test whether the logprob for ```pt.max``` or ```pt.min``` for non i.i.d is correctly rejected""" - x = pm.Normal.dist([0, 1, 2, 3, 4], 1, shape=(5,)) - x.name = "x" - x_m = pt_op(x, axis=-1) - x_m_value = pt.scalar("x_value") - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): - logp(x_m, x_m_value) - - -@pytest.mark.parametrize( - "pt_op", - [ - pt.max, - pt.min, - ], -) -def test_multivariate_rv_fails(pt_op): - _alpha = pt.scalar() - _k = pt.iscalar() - x = pm.StickBreakingWeights.dist(_alpha, _k) - x.name = "x" - x_m = pt_op(x, axis=-1) - x_m_value = pt.scalar("x_value") - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): - logp(x_m, x_m_value) - - -@pytest.mark.parametrize( - "pt_op", - [ - pt.max, - pt.min, - ], -) -def test_categorical(pt_op): - """Test whether the logprob for ```pt.max``` for unsupported distributions is correctly rejected""" - x = pm.Categorical.dist([1, 1, 1, 1], shape=(5,)) - x.name = "x" - x_m = pt_op(x, axis=-1) - x_m_value = pt.scalar("x_value", dtype=x.type.dtype) - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): - logp(x_m, x_m_value) - - -@pytest.mark.parametrize( - "pt_op", - [ - pt.max, - pt.min, - ], -) -def test_non_supp_axis(pt_op): - """Test whether the logprob for ```pt.max``` for unsupported axis is correctly rejected""" - x = pt.random.normal(0, 1, size=(3, 3)) - x.name = "x" - x_m = pt_op(x, axis=-1) - x_m_value = pt.vector("x_value") - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): - x_max_logprob = logp(x_m, x_m_value) - - -@pytest.mark.parametrize( - "shape, value, axis", - [ - (3, 0.85, -1), - (3, 0.01, 0), - (2, 0.2, None), - (4, 0.5, 0), - ((3, 4), 0.9, None), - ((3, 4), 0.75, (1, 0)), - ], -) -def test_max_logprob(shape, value, axis): - """Test whether the logprob for ```pt.max``` produces the corrected - - The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here: - U_1, \\dots, U_n \\stackrel{\text{i.i.d.}}{\\sim} \text{Uniform}(0, 1) \\Rightarrow U_{(k)} \\sim \text{Beta}(k, n + 1- k) - for all 1<=k<=n - """ - x = pt.random.uniform(0, 1, size=shape) - x.name = "x" - x_max = pt.max(x, axis=axis) - x_max_value = pt.scalar("x_max_value") - x_max_logprob = logp(x_max, x_max_value) - - assert_no_rvs(x_max_logprob) - - test_value = value - - n = np.prod(shape) - beta_rv = pt.random.beta(n, 1, name="beta") - beta_vv = beta_rv.clone() - beta_rv_logprob = logp(beta_rv, beta_vv) - - np.testing.assert_allclose( - beta_rv_logprob.eval({beta_vv: test_value}), - (x_max_logprob.eval({x_max_value: test_value})), - rtol=1e-06, +class TestMax: + def test_argmax(self): + """Test whether the logprob for ```pt.argmax``` is correctly rejected""" + x = pt.random.normal(0, 1, size=(3,)) + x.name = "x" + x_argmax = pt.argmax(x, axis=-1) + x_max_value = pt.scalar("x_max_value", dtype=x_argmax.type.dtype) + + with pytest.raises( + RuntimeError, match=re.escape("Logprob method not implemented for Argmax") + ): + logp(x_argmax, x_max_value) + + @pytest.mark.parametrize( + "pt_op", + [ + pt.max, + pt.min, + ], ) - - -@pytest.mark.parametrize( - "shape, value, axis", - [ - (3, 0.85, -1), - (3, 0.01, 0), - (2, 0.2, None), - (4, 0.5, 0), - ((3, 4), 0.9, None), - ((3, 4), 0.75, (1, 0)), - ], -) -def test_min_logprob(shape, value, axis): - """Test whether the logprob for ```pt.mix``` produces the corrected - The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here: - U_1, \\dots, U_n \\stackrel{\text{i.i.d.}}{\\sim} \text{Uniform}(0, 1) \\Rightarrow U_{(k)} \\sim \text{Beta}(k, n + 1- k) - for all 1<=k<=n - """ - x = pt.random.uniform(0, 1, size=shape) - x.name = "x" - x_min = pt.min(x, axis=axis) - x_min_value = pt.scalar("x_min_value") - x_min_logprob = logp(x_min, x_min_value) - - assert_no_rvs(x_min_logprob) - - test_value = value - - n = np.prod(shape) - beta_rv = pt.random.beta(1, n, name="beta") - beta_vv = beta_rv.clone() - beta_rv_logprob = logp(beta_rv, beta_vv) - - np.testing.assert_allclose( - beta_rv_logprob.eval({beta_vv: test_value}), - (x_min_logprob.eval({x_min_value: test_value})), - rtol=1e-06, + def test_non_iid_fails(self, pt_op): + """Test whether the logprob for ```pt.max``` or ```pt.min``` for non i.i.d is correctly rejected""" + x = pt.random.normal([0, 1, 2, 3, 4], 1, size=(5,)) + x.name = "x" + x_m = pt_op(x, axis=-1) + x_m_value = pt.scalar("x_value") + with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): + logp(x_m, x_m_value) + + @pytest.mark.parametrize( + "pt_op", + [ + pt.max, + pt.min, + ], ) - - -def test_min_non_mul_elemwise_fails(): - """Test whether the logprob for ```pt.min``` for non-mul elemwise RVs is rejected correctly""" - x = pt.log(pt.random.beta(0, 1, size=(3,))) - x.name = "x" - x_min = pt.min(x, axis=-1) - x_min_value = pt.scalar("x_min_value") - with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): - logp(x_min, x_min_value) - - -@pytest.mark.parametrize( - "mu, size, value, axis", - [(2, 3, 1, -1), (2, 3, 1, 0), (1, 2, 2, None), (0, 4, 0, 0)], -) -def test_max_discrete(mu, size, value, axis): - x = pm.Poisson.dist(name="x", mu=mu, size=size) - x_max = pt.max(x, axis=axis) - x_max_value = pt.scalar("x_max_value", dtype=x.type.dtype) - x_max_logprob = logp(x_max, x_max_value) - - test_value = value - - n = size - exp_rv = sp.poisson(mu).cdf(test_value) ** n - exp_rv_prev = sp.poisson(mu).cdf(test_value - 1) ** n - - np.testing.assert_allclose( - np.log(exp_rv - exp_rv_prev), - (x_max_logprob.eval({x_max_value: test_value})), - rtol=1e-06, + def test_multivariate_rv_fails(self, pt_op): + _alpha = pt.vector() + x = pt.random.dirichlet(_alpha) + x.name = "x" + x_m = pt_op(x, axis=-1) + x_m_value = pt.scalar("x_value") + with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): + logp(x_m, x_m_value) + + @pytest.mark.parametrize( + "pt_op", + [ + pt.max, + pt.min, + ], ) - - -@pytest.mark.parametrize( - "mu, n, test_value, axis", - [(2, 3, 1, -1), (2, 3, 1, 0), (1, 2, 2, None), (0, 4, 0, 0)], -) -def test_min_discrete(mu, n, test_value, axis): - x = pm.Poisson.dist(name="x", mu=mu, size=(n,)) - x_min = pt.min(x, axis=axis) - x_min_value = pt.scalar("x_min_value", dtype=x.type.dtype) - x_min_logprob = logp(x_min, x_min_value) - - sf_before = 1 - sp.poisson(mu).cdf(test_value - 1) - sf = 1 - sp.poisson(mu).cdf(test_value) - - expected_logp = np.log(sf_before**n - sf**n) - - np.testing.assert_allclose( - x_min_logprob.eval({x_min_value: test_value}), - expected_logp, - rtol=1e-06, + def test_categorical(self, pt_op): + """Test whether the logprob for ```pt.max``` for unsupported distributions is correctly rejected""" + x = pt.random.categorical([1, 1, 1, 1], size=(5,)) + x.name = "x" + x_m = pt_op(x, axis=-1) + x_m_value = pt.scalar("x_value", dtype=x.type.dtype) + with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): + logp(x_m, x_m_value) + + @pytest.mark.parametrize( + "pt_op", + [ + pt.max, + pt.min, + ], ) - - -def test_min_max_bernoulli(): - p = 0.7 - q = 1 - p - n = 3 - x = pm.Bernoulli.dist(name="x", p=p, shape=(n,)) - value = pt.scalar("value", dtype=int) - - max_logp_fn = pytensor.function([value], pm.logp(pt.max(x), value)) - np.testing.assert_allclose(max_logp_fn(0), np.log(q**n)) - np.testing.assert_allclose(max_logp_fn(1), np.log(1 - q**n)) - - min_logp_fn = pytensor.function([value], pm.logp(pt.min(x), value)) - np.testing.assert_allclose(min_logp_fn(1), np.log(p**n)) - np.testing.assert_allclose(min_logp_fn(0), np.log(1 - p**n)) - - -def test_non_measurable_max_grad(): - # Regression test for https://github.com/pymc-devs/pytensor/issues/711 - x = pt.random.normal(0, 1, size=(3,)) - max_x = x.max() - y = pt.random.normal(max_x, 1) - - x_vv = x.type() - y_vv = y.type() - logp_terms = conditional_logp({x: x_vv, y: y_vv}).values() - joint_logp = pt.sum([term.sum() for term in logp_terms]) - - # Test that calling gradient does not raise a NotImplementedError - pt.grad(joint_logp, x_vv) + def test_non_supp_axis(self, pt_op): + """Test whether the logprob for ```pt.max``` for unsupported axis is correctly rejected""" + x = pt.random.normal(0, 1, size=(3, 3)) + x.name = "x" + x_m = pt_op(x, axis=-1) + x_m_value = pt.vector("x_value") + with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): + x_max_logprob = logp(x_m, x_m_value) + + @pytest.mark.parametrize( + "shape, value, axis", + [ + (3, 0.85, -1), + (3, 0.01, 0), + (2, 0.2, None), + (4, 0.5, 0), + ((3, 4), 0.9, None), + ((3, 4), 0.75, (1, 0)), + ], + ) + def test_max_logprob(self, shape, value, axis): + """Test whether the logprob for ```pt.max``` produces the corrected + + The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here: + U_1, \\dots, U_n \\stackrel{\text{i.i.d.}}{\\sim} \text{Uniform}(0, 1) \\Rightarrow U_{(k)} \\sim \text{Beta}(k, n + 1- k) + for all 1<=k<=n + """ + x = pt.random.uniform(0, 1, size=shape) + x.name = "x" + x_max = pt.max(x, axis=axis) + x_max_value = pt.scalar("x_max_value") + x_max_logprob = logp(x_max, x_max_value) + + assert_no_rvs(x_max_logprob) + + test_value = value + + n = np.prod(shape) + beta_rv = pt.random.beta(n, 1, name="beta") + beta_vv = beta_rv.clone() + beta_rv_logprob = logp(beta_rv, beta_vv) + + np.testing.assert_allclose( + beta_rv_logprob.eval({beta_vv: test_value}), + (x_max_logprob.eval({x_max_value: test_value})), + rtol=1e-06, + ) + + @pytest.mark.parametrize( + "shape, value, axis", + [ + (3, 0.85, -1), + (3, 0.01, 0), + (2, 0.2, None), + (4, 0.5, 0), + ((3, 4), 0.9, None), + ((3, 4), 0.75, (1, 0)), + ], + ) + def test_min_logprob(self, shape, value, axis): + """Test whether the logprob for ```pt.mix``` produces the corrected + The fact that order statistics of i.i.d. uniform RVs ~ Beta is used here: + U_1, \\dots, U_n \\stackrel{\text{i.i.d.}}{\\sim} \text{Uniform}(0, 1) \\Rightarrow U_{(k)} \\sim \text{Beta}(k, n + 1- k) + for all 1<=k<=n + """ + x = pt.random.uniform(0, 1, size=shape) + x.name = "x" + x_min = pt.min(x, axis=axis) + x_min_value = pt.scalar("x_min_value") + x_min_logprob = logp(x_min, x_min_value) + + assert_no_rvs(x_min_logprob) + + test_value = value + + n = np.prod(shape) + beta_rv = pt.random.beta(1, n, name="beta") + beta_vv = beta_rv.clone() + beta_rv_logprob = logp(beta_rv, beta_vv) + + np.testing.assert_allclose( + beta_rv_logprob.eval({beta_vv: test_value}), + (x_min_logprob.eval({x_min_value: test_value})), + rtol=1e-06, + ) + + def test_min_non_mul_elemwise_fails(self): + """Test whether the logprob for ```pt.min``` for non-mul elemwise RVs is rejected correctly""" + x = pt.log(pt.random.beta(0, 1, size=(3,))) + x.name = "x" + x_min = pt.min(x, axis=-1) + x_min_value = pt.scalar("x_min_value") + with pytest.raises(RuntimeError, match=re.escape("Logprob method not implemented")): + logp(x_min, x_min_value) + + @pytest.mark.parametrize( + "lam, size, value, axis", + [(2, 3, 1, -1), (2, 3, 1, 0), (1, 2, 2, None), (0, 4, 0, 0)], + ) + def test_max_discrete(self, lam, size, value, axis): + x = pt.random.poisson(name="x", lam=lam, size=size) + x_max = pt.max(x, axis=axis) + x_max_value = pt.scalar("x_max_value", dtype=x.type.dtype) + x_max_logprob = logp(x_max, x_max_value) + + test_value = value + + n = size + exp_rv = sp.poisson(lam).cdf(test_value) ** n + exp_rv_prev = sp.poisson(lam).cdf(test_value - 1) ** n + + np.testing.assert_allclose( + np.log(exp_rv - exp_rv_prev), + (x_max_logprob.eval({x_max_value: test_value})), + rtol=1e-06, + ) + + @pytest.mark.parametrize( + "lam, n, test_value, axis", + [(2, 3, 1, -1), (2, 3, 1, 0), (1, 2, 2, None), (0, 4, 0, 0)], + ) + def test_min_discrete(self, lam, n, test_value, axis): + x = pt.random.poisson(name="x", lam=lam, size=(n,)) + x_min = pt.min(x, axis=axis) + x_min_value = pt.scalar("x_min_value", dtype=x.type.dtype) + x_min_logprob = logp(x_min, x_min_value) + + sf_before = 1 - sp.poisson(lam).cdf(test_value - 1) + sf = 1 - sp.poisson(lam).cdf(test_value) + + expected_logp = np.log(sf_before**n - sf**n) + + np.testing.assert_allclose( + x_min_logprob.eval({x_min_value: test_value}), + expected_logp, + rtol=1e-06, + ) + + def test_min_max_bernoulli(self): + p = 0.7 + q = 1 - p + n = 3 + x = pt.random.bernoulli(name="x", p=p, size=(n,)) + value = pt.scalar("value", dtype=int) + + max_logp_fn = pytensor.function([value], logp(pt.max(x), value)) + np.testing.assert_allclose(max_logp_fn(0), np.log(q**n)) + np.testing.assert_allclose(max_logp_fn(1), np.log(1 - q**n)) + + min_logp_fn = pytensor.function([value], logp(pt.min(x), value)) + np.testing.assert_allclose(min_logp_fn(1), np.log(p**n)) + np.testing.assert_allclose(min_logp_fn(0), np.log(1 - p**n)) + + def test_non_measurable_max_grad(self): + # Regression test for https://github.com/pymc-devs/pytensor/issues/711 + x = pt.random.normal(0, 1, size=(3,)) + max_x = x.max() + y = pt.random.normal(max_x, 1) + + x_vv = x.type() + y_vv = y.type() + logp_terms = conditional_logp({x: x_vv, y: y_vv}).values() + joint_logp = pt.sum([term.sum() for term in logp_terms]) + + # Test that calling gradient does not raise a NotImplementedError + pt.grad(joint_logp, x_vv) From 07178763a10c6a503a8a5f15f7d9adf274dbb4f9 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sun, 25 Jan 2026 15:45:17 +0100 Subject: [PATCH 4/4] Derive categorical from argmin/argmax operations --- pymc/logprob/order.py | 187 +++++++++++++++++++++++++++++++++++- tests/logprob/test_order.py | 169 ++++++++++++++++++++++++++++++++ 2 files changed, 355 insertions(+), 1 deletion(-) diff --git a/pymc/logprob/order.py b/pymc/logprob/order.py index 30f6a32565..af6ef9c486 100644 --- a/pymc/logprob/order.py +++ b/pymc/logprob/order.py @@ -35,14 +35,25 @@ # SOFTWARE. from typing import cast +import numpy as np import pytensor.tensor as pt +from numpy.lib.array_utils import normalize_axis_tuple +from pytensor.graph import ancestors from pytensor.graph.basic import Apply from pytensor.graph.fg import FunctionGraph from pytensor.graph.rewriting.basic import node_rewriter -from pytensor.tensor.math import Max +from pytensor.scalar import Add, Mul +from pytensor.tensor import get_underlying_scalar_constant_value +from pytensor.tensor.elemwise import Elemwise +from pytensor.tensor.exceptions import NotScalarConstantError +from pytensor.tensor.math import Argmax, Max, variadic_add, variadic_mul +from pytensor.tensor.random.basic import ExponentialRV, GumbelRV +from pytensor.tensor.rewriting.basic import broadcasted_by +from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable +from pymc.distributions.continuous import WeibullBetaRV from pymc.logprob.abstract import ( MeasurableElemwise, MeasurableOp, @@ -122,6 +133,7 @@ def find_measurable_max(fgraph: FunctionGraph, node: Apply) -> list[TensorVariab "find_measurable_max", find_measurable_max, "basic", + "order", "max", ) @@ -158,3 +170,176 @@ def max_logprob_discrete(op, values, base_rv, **kwargs): n = pt.prod(base_rv_shape) return logdiffexp(n * logcdf, n * logcdf_prev) + + +@node_rewriter([ExponentialRV, GumbelRV, WeibullBetaRV]) +def lift_loc_scale(fgraph, node): + """Rewrite rv(loc, scale) * s + l as rv(loc * s + l, scale * s), when they admit such parametrization. + + Currently, this rewrite targets just RVs accepted by categorical_from_argmax, but can be generalized beyond this use case. + """ + rv = node.out + clients = fgraph.clients[rv] + if len(clients) != 1: + return None + [(add_mul_node, idx)] = clients + + if not ( + isinstance(add_mul_node.op, Elemwise) and isinstance(add_mul_node.op.scalar_op, Add | Mul) + ): + return None + + match node.op: + case ExponentialRV(): + loc = 0 + rng, size, scale = node.inputs + case WeibullBetaRV(): + loc = 0 + rng, size, shape, scale = node.inputs + case GumbelRV(): + rng, size, loc, scale = node.inputs + case _: + raise NotImplementedError(f"Unexpected op {node.op}") + + negative_scale = False + if isinstance(add_mul_node.op.scalar_op, Add): + if not isinstance(node.op, GumbelRV): + # Only Gumbel allows lifting loc + return None + loc += variadic_add(*(t for i, t in enumerate(add_mul_node.inputs) if i != idx)) + else: + extra_scale = variadic_mul(*(t for i, t in enumerate(add_mul_node.inputs) if i != idx)) + # The rewrite is only valid if the scale non-negative. + try: + [extra_scale_const] = constant_fold([extra_scale]) + except NotScalarConstantError: + # Not-constant scale. Get out + # TODO: Allow more cases when we have standard machinery to infer sign of symbolic operations + return None + unique_sign = np.unique(np.sign(extra_scale_const)) + if unique_sign.size == 2: + # There is mixed sign in the scale (or zero). Get out + return None + negative_scale = unique_sign == -1 + if negative_scale: + if (extra_scale_const == -1).all(): + # There is no scale to lift, it's just a negated rv (happens in argmin(rv)) + return None + # Scale is homogenously negative, make it positive, and return rv * -1 later so other rewrites can handle it + extra_scale *= -1 + loc *= extra_scale + scale *= extra_scale + + # We can't lift the argument if either + # 1. the argument it's broadcasting the RV + # 2. the RV shows up in the lifted args (as in rv + rv or rv * rv) + # We check with loc as that is altered in both branches + if broadcasted_by(rv, loc) or rv in ancestors([loc]): + return None + + match node.op: + case ExponentialRV(): + lifted_rv = node.op.make_node(rng, size, scale).out + case WeibullBetaRV(): + # WeibullBetaRV is a SymbolicRandomVariable, we can't simply pass arguments of a different type + lifted_rv = WeibullBetaRV.rv_op(shape, scale, rng=rng, size=size) + case GumbelRV(): + lifted_rv = node.op.make_node(rng, size, loc, scale).out + + if negative_scale: + lifted_rv *= -1 + + return {add_mul_node.out: lifted_rv} + + +@node_rewriter([Argmax]) +def categorical_from_argmax(fgraph, node): + """Convert closed from argmax/argmin to equivalent categorical.""" + + def is_minus_1(x): + try: + return get_underlying_scalar_constant_value(x, max_recur=3) == -1 + except NotScalarConstantError: + return False + + [base_var] = node.inputs + base_node = base_var.owner + + if base_node is None: + return None + if not filter_measurable_variables([base_var]): + return None + + argmax_axes = node.op.axis + if argmax_axes is None: + argmax_axes = tuple(range(base_var.ndim)) + else: + argmax_axes = normalize_axis_tuple(argmax_axes, base_var.ndim) + + probs = None + + if isinstance(base_node.op, GumbelRV): + # argmax(gumbel(loc, scale)) -> categorical(exp(loc / scale)) + rng, size, loc, scale = base_node.inputs + + # gumbel scale has to be constant across the argmax axes + if not all(b for i, b in enumerate(scale.type.broadcastable) if i in argmax_axes): + return None + + # Dividing by scale also broadcasts probs to the size implied by (loc, scale) + probs = pt.exp(loc / scale) + + # Check if we have an Argmin + # Argmin is internally represented as Argmax(-x), and -x is canonicalized as x * -1 + elif ( + len(base_node.inputs) == 2 + and isinstance(base_node.op, Elemwise) + and isinstance(base_node.op.scalar_op, Mul) + and ((right_neg := is_minus_1(base_node.inputs[1])) or is_minus_1(base_node.inputs[0])) + ): + base_var = base_node.inputs[0 if right_neg else 1] + base_node = base_var.owner + if base_node is None: + return None + + if isinstance(base_node.op, ExponentialRV): + # argmin(exponential(rate)) -> categorical(rate) + rng, size, scale = base_var.owner.inputs + probs = 1 / scale + + elif isinstance(base_node.op, WeibullBetaRV): + # argmin(weibull(shape, scale)) -> categorical(scale ** -shape / Σ(scale ** -shape)) + rng, size, shape, scale = base_node.inputs + + # weibull shape has to be constant across the argmin axes + if not all(b for i, b in enumerate(shape.type.broadcastable) if i in argmax_axes): + return None + + probs = scale**-shape + + if probs is None: + return None + + if not isinstance(size.type, NoneTypeT): + # Make probs explicit to facilitate logic below + probs = pt.broadcast_to(probs, size) + + # Join axes probs at the last axis (core axis of Categorical) + n_axes = len(argmax_axes) + probs = pt.moveaxis(probs, argmax_axes, tuple(range(-n_axes, 0))) + probs = pt.join_dims(probs, -n_axes, n_axes) + + # Normalize probs and create categorical + probs /= probs.sum(-1, keepdims=True) + return [pt.random.categorical(probs, rng=rng, size=None)] + + +measurable_ir_rewrites_db.register("lift_loc_scale", lift_loc_scale, "basic", "lift_rv_args") + +measurable_ir_rewrites_db.register( + "categorical_from_argmax", + categorical_from_argmax, + "basic", + "order", + "argmax", +) diff --git a/tests/logprob/test_order.py b/tests/logprob/test_order.py index d7300b18fc..4760146528 100644 --- a/tests/logprob/test_order.py +++ b/tests/logprob/test_order.py @@ -42,6 +42,7 @@ import pytest import scipy.stats as sp +from pymc.distributions import Weibull from pymc.logprob import conditional_logp, logp from pymc.testing import assert_no_rvs @@ -280,3 +281,171 @@ def test_non_measurable_max_grad(self): # Test that calling gradient does not raise a NotImplementedError pt.grad(joint_logp, x_vv) + + +class TestCategoricalFromArgmax: + # np.pi, np.e, are used throughout to emphasize they don't matter in the end + + def test_exponential(self): + axis = 1 + size = (2, 3, 4) + probs = np.array([0.1, 0.3, 0.6]) + scale = 1 / (probs * np.e)[None, :, None] + + s = pt.random.exponential(scale=1, size=size) * scale + x = pt.argmin(s, axis=axis) + x_vv = x.type() + x_prob = pt.exp(logp(x, x_vv)) + + x_vv_test = [[0, 1, 2, 0], [1, 2, 2, 2]] + expected_probs = probs[x_vv_test] + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + expected_probs, + ) + + # Test same result for lifted scale + equiv_s = pt.random.exponential(scale=scale, size=size) + equiv_x = pt.argmin(equiv_s, axis=axis) + equiv_x_prob = pt.exp(logp(equiv_x, x_vv)) + np.testing.assert_allclose(equiv_x_prob.eval({x_vv: x_vv_test}), expected_probs) + + # Test argmax not supported + invalid_x = pt.argmax(s, axis=axis) + with pytest.raises(NotImplementedError): + logp(invalid_x, x_vv) + + def test_weibull(self): + axis = 1 + size = (2, 3, 4) + probs = np.array([0.1, 0.3, 0.6]) + # shape is IID on the argmax axis + iid_shape = np.array([np.pi, np.e])[:, None, None] + # probs = scale ** -shape, normalized + scale = probs[None, :, None] ** (-1 / iid_shape) + + s = Weibull.dist(alpha=iid_shape, beta=1, size=size) * scale + x = pt.argmin(s, axis=axis) + x_vv = x.type() + x_prob = pt.exp(logp(x, x_vv)) + + x_vv_test = [[0, 1, 2, 0], [1, 2, 2, 2]] + expected_probs = probs[x_vv_test] + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + expected_probs, + ) + + # Test argmax not supported + invalid_x = pt.argmax(s, axis=axis) + with pytest.raises(NotImplementedError): + logp(invalid_x, x_vv) + + # Test non-iid shape on argmin axis not supported + non_iid_shape = np.array([1, np.pi, np.e])[None, :, None] + invalid_s = Weibull.dist(alpha=non_iid_shape, beta=1, size=size) * scale + invalid_x = pt.argmin(invalid_s, axis=axis) + with pytest.raises(NotImplementedError): + logp(invalid_x, x_vv) + + def test_gumbel(self): + axis = -2 + size = (2, 3, 4) + probs = np.array([0.1, 0.3, 0.6]) + # scale is IID on the argmax axis + iid_scale = np.array([np.pi, np.e])[:, None, None] + # loc is proportional to the scale, so that probs = softmax(loc / scale) don't depend on it + loc = pt.log(probs * np.e)[None, :, None] * iid_scale + + s = pt.random.gumbel(loc=0, scale=1, size=size) * iid_scale + loc + x = pt.argmax(s, axis=axis) + x_vv = x.type() + x_prob = pt.exp(logp(x, x_vv)) + + x_vv_test = [[0, 1, 2, 0], [1, 2, 2, 2]] + expected_probs = probs[x_vv_test] + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + expected_probs, + ) + + # Test same result for lifted loc / scale + equiv_s = pt.random.gumbel(loc=loc, scale=iid_scale, size=size) + equiv_x = pt.argmax(equiv_s, axis=axis) + equiv_x_prob = pt.exp(logp(equiv_x, x_vv)) + np.testing.assert_allclose(equiv_x_prob.eval({x_vv: x_vv_test}), expected_probs) + + # Test argmin not supported + invalid_x = pt.argmin(s, axis=axis) + with pytest.raises(NotImplementedError): + logp(invalid_x, x_vv) + + # Test non-iid scale on argmax axis not supported + non_iid_scale = np.array([1, np.pi, np.e])[None, :, None] + invalid_s = pt.random.gumbel( + loc=loc, + scale=non_iid_scale, + size=size, + ) + invalid_x = pt.argmax(invalid_s, axis=axis) + with pytest.raises(NotImplementedError): + logp(invalid_x, x_vv) + + def test_implicit_size(self): + probs = np.array([[0.1, 0.5], [0.9, 0.5]]) + s = pt.random.exponential(scale=1 / probs) + x = pt.argmin(s, axis=0) + + x_vv = x.type() + x_prob = pt.exp(logp(x, x_vv)) + + x_vv_test = [0, 0] + np.testing.assert_allclose(x_prob.eval({x_vv: x_vv_test}), [0.1, 0.5]) + + def test_multiple_axes(self): + probs = np.array([0.1, 0.2, 0.7]) + s = pt.random.exponential(scale=1 / probs, size=(2, 4, 3)) + x = pt.argmin(s, axis=(0, -1)) + assert x.type.shape == (4,) + x_vv = x.type() + + x_prob = pt.exp(logp(x, x_vv)) + x_vv_test = [0, 1, 0, 5] + # We are selecting among 2x repeated scales + expected = (np.tile(probs, 2) / 2)[x_vv_test] + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + expected, + ) + + x = pt.argmin(s, axis=None) + assert x.type.shape == () + x_vv = x.type() + x_prob = pt.exp(logp(x, x_vv)) + x_vv_test = 15 + # We are selecting among 8x repeated scales + expected = (np.broadcast_to(probs / 8, (2, 4, 3)).ravel()[x_vv_test],) + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + expected, + ) + + @pytest.mark.xfail( + raises=NotImplementedError, + reason="Implicitly depends on https://github.com/pymc-devs/pytensor/issues/1851", + ) + def test_derived_frechet(self): + probs = np.array([0.1, 0.2, 0.7]) + # Frechet is simply an exponentiated gumbel + # Since exponentiation is monotonic it doesn't change the result of argmax, + # and the probability is therefore the same as that of the argmin(gumbel) + s = pt.exp(pt.random.gumbel(loc=pt.log(probs))) + x = pt.argmax(s) + x_vv = x.type() + + x_prob = pt.exp(logp(x, x_vv)) + x_vv_test = 1 + np.testing.assert_allclose( + x_prob.eval({x_vv: x_vv_test}), + probs[x_vv_test], + )