From 1ea33552832fa48c6489b6fee981f19b13e879ad Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 30 Jun 2023 11:20:40 +0200 Subject: [PATCH] Derive probability for broadcast operation Broadcasting lowers to Alloc; the logp scores the base variable at the representative entries and returns -inf where the value is inconsistent with a broadcast, elementwise over the base's batch dimensions. The expanded and broadcast dimensions are degenerate copies of the base entries, so they are consumed like support dimensions and disappear from the logp. Without meta information about the support axes (#6360), other rewrites would treat the copies as independent entries (e.g., counting the jacobian of a transform once per copy), so the broadcast is only claimed when directly valued. For the same reason the implicit broadcasting of measurable inputs inside elemwise transforms is not attempted. --- .github/workflows/tests.yml | 1 + pymc/logprob/__init__.py | 1 + pymc/logprob/shape.py | 137 +++++++++++++++++++++++++++++++++++ tests/logprob/test_shape.py | 95 ++++++++++++++++++++++++ tests/logprob/test_tensor.py | 3 +- 5 files changed, 235 insertions(+), 2 deletions(-) create mode 100644 pymc/logprob/shape.py create mode 100644 tests/logprob/test_shape.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a4169733ff..f4e889e97a 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -135,6 +135,7 @@ jobs: tests/logprob/test_order.py tests/logprob/test_rewriting.py tests/logprob/test_scan.py + tests/logprob/test_shape.py tests/logprob/test_tensor.py tests/logprob/test_switch.py tests/logprob/test_transform_value.py diff --git a/pymc/logprob/__init__.py b/pymc/logprob/__init__.py index 985f9f489d..a48e1ca5ac 100644 --- a/pymc/logprob/__init__.py +++ b/pymc/logprob/__init__.py @@ -55,6 +55,7 @@ import pymc.logprob.mixture import pymc.logprob.order import pymc.logprob.scan +import pymc.logprob.shape import pymc.logprob.switch import pymc.logprob.tensor import pymc.logprob.transforms diff --git a/pymc/logprob/shape.py b/pymc/logprob/shape.py new file mode 100644 index 0000000000..58b35e6c38 --- /dev/null +++ b/pymc/logprob/shape.py @@ -0,0 +1,137 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytensor.tensor as pt + +from pytensor.graph.rewriting.basic import node_rewriter +from pytensor.tensor.basic import Alloc + +from pymc.logprob.abstract import MeasurableOp, _logprob, _logprob_helper +from pymc.logprob.rewriting import measurable_ir_rewrites_db +from pymc.logprob.utils import ( + check_potential_measurability, + filter_measurable_variables, + get_related_valued_nodes, +) + + +class MeasurableBroadcast(MeasurableOp, Alloc): + """A placeholder used to specify a log-likelihood for a broadcast sub-graph.""" + + +@_logprob.register(MeasurableBroadcast) +def broadcast_logprob(op, values, rv, *shape, **kwargs): + """Log-probability expression for (statically-)broadcasted RV. + + The probability is the same as the base RV, if no broadcasting had happened. + The broadcast dimensions are degenerate copies of the base entries, so they are + consumed like support dimensions and disappear from the logp: + + ``logp(broadcast_to(normal(size=(3, 1)), (2, 3, 4)), zeros((2, 3, 4))) == logp(normal(size=(3,)), zeros((3,)))`` + + And zero if the value couldn't have possibly originated via broadcasting: + + ``logp(broadcast_to(normal(size=(1,)), (3,)), [1, 2, 3]) == -np.inf`` + + The consistency check is elementwise over the base variable's batch dimensions, + so entries that were not broadcast from each other keep their own logp. + """ + [value] = values + + n_new_dims = len(shape) - rv.ndim + assert n_new_dims >= 0 + + # Enumerate broadcasted dims + expanded_dims = tuple(range(n_new_dims)) + broadcast_dims = tuple( + i + n_new_dims + for i, (v_bcast, rv_bcast) in enumerate( + zip(value.broadcastable[n_new_dims:], rv.broadcastable) + ) + if (not v_bcast) and rv_bcast + ) + + # "Unbroadcast" value via indexing. + # All entries in the broadcasted dimensions should be the same, so we simply select + # the first of each. Broadcast dims are re-inserted with expand_dims (rather than + # sliced with `0:1`) so they are statically known to be broadcastable. + indices = [] + for i in range(value.ndim): + # Remove expanded and broadcasted (but not expanded) dims + if i in expanded_dims or i in broadcast_dims: + indices.append(0) + else: + indices.append(slice(None)) + + unbroadcast_value = value[tuple(indices)] + # The base variable still carries the broadcast dims (with length 1); they are + # re-inserted with expand_dims so they are statically known to be broadcastable + rv_value = unbroadcast_value + if broadcast_dims: + rv_value = pt.expand_dims(rv_value, tuple(d - n_new_dims for d in broadcast_dims)) + logp = _logprob_helper(rv, rv_value) + + # The broadcast dims are consumed like support dims and disappear from the logp + core_ndim = rv_value.ndim - logp.ndim + squeeze_axes = tuple(d - n_new_dims for d in broadcast_dims if (d - n_new_dims) < logp.ndim) + if squeeze_axes: + logp = pt.squeeze(logp, axis=squeeze_axes) + + # Check that dependent values were indeed identical, by comparing with a re-broadcasted value. + # The check is reduced only over the expanded/broadcast dimensions (and any support + # dimensions consumed by the base logp), so unrelated batch entries do not + # contaminate each other. + # Note: This could fail due to float-precision issues. + # If that proves to be a problem we should switch to `pt.allclose` + valid_value = pt.broadcast_to(rv_value, shape) + core_dims = tuple(range(value.ndim - core_ndim, value.ndim)) + reduced_dims = tuple(sorted({*expanded_dims, *broadcast_dims, *core_dims})) + check = pt.all(pt.eq(value, valid_value), axis=reduced_dims) + + return pt.switch(check, logp, -np.inf) + + +@node_rewriter([Alloc]) +def find_measurable_broadcast(fgraph, node): + r"""Find measurable broadcasts ``broadcast_to(rv, shape)``.""" + if isinstance(node.op, MeasurableOp): + return None + + base_rv, *shape = node.inputs + + if not filter_measurable_variables([base_rv]): + return None + + if check_potential_measurability(shape): + return None + + # The broadcast dimensions are degenerate copies of the base variable's entries. + # Without meta information about the support axes, other rewrites would treat the + # copies as independent entries (e.g., counting the jacobian of a transform once + # per copy), so the broadcast is only claimed when directly valued, where no other + # rewrite needs to reason about the returned logp. + # TODO: When we include the support axis as meta information in each intermediate + # MeasurableVariable, we can lift this restriction (see https://github.com/pymc-devs/pymc/issues/6360) + if not any(get_related_valued_nodes(fgraph, node)): + return None + + return [MeasurableBroadcast()(base_rv, *shape)] + + +measurable_ir_rewrites_db.register( + "find_measurable_broadcast", + find_measurable_broadcast, + "basic", + "shape", +) diff --git a/tests/logprob/test_shape.py b/tests/logprob/test_shape.py new file mode 100644 index 0000000000..bf304fd7c1 --- /dev/null +++ b/tests/logprob/test_shape.py @@ -0,0 +1,95 @@ +# Copyright 2024 - present The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import pytensor +import pytensor.tensor as pt +import pytest +import scipy.stats as st + +from pymc import logp + + +def test_measurable_broadcast(): + b_shape = pt.vector("b_shape", shape=(3,), dtype=int) + + x = pt.random.normal(size=(3, 1)) + bcast_x = pt.broadcast_to(x, shape=b_shape) + bcast_x.name = "bcast_x" + + bcast_x_value = bcast_x.clone() + logp_bcast_x = logp(bcast_x, bcast_x_value) + logp_fn = pytensor.function([b_shape, bcast_x_value], logp_bcast_x, on_unused_input="ignore") + + # The expanded and broadcast dimensions are consumed like support dimensions: + # the logp has the base variable's remaining batch shape + # (assert_allclose also asserts shapes match, if neither is scalar) + np.testing.assert_allclose( + logp_fn([1, 3, 1], np.zeros((1, 3, 1))), + st.norm.logpdf(np.zeros(3)), + ) + np.testing.assert_allclose( + logp_fn([1, 3, 5], np.zeros((1, 3, 5))), + st.norm.logpdf(np.zeros(3)), + ) + np.testing.assert_allclose( + logp_fn([2, 3, 5], np.broadcast_to(np.arange(3).reshape(1, 3, 1), (2, 3, 5))), + st.norm.logpdf(np.arange(3)), + ) + # Invalid broadcast value + np.testing.assert_array_equal( + logp_fn([1, 3, 5], np.arange(3 * 5).reshape(1, 3, 5)), + np.full(shape=(3,), fill_value=-np.inf), + ) + # The invalidity check is elementwise over the base batch dimensions: an + # inconsistent row only invalidates its own logp + partially_valid = np.broadcast_to(np.arange(3).reshape(1, 3, 1), (1, 3, 5)).copy() + partially_valid[0, 1, 3] = 99.0 + np.testing.assert_allclose( + logp_fn([1, 3, 5], partially_valid), + np.where([True, False, True], st.norm.logpdf(np.arange(3)), -np.inf), + ) + + +def test_measurable_broadcast_multivariate(): + x = pt.random.dirichlet(pt.ones(3), size=(1,)) + bcast_x = pt.broadcast_to(x, (5, 3)) + + bcast_x_value = bcast_x.clone() + logp_bcast_x = logp(bcast_x, bcast_x_value) + + rng = np.random.default_rng(170) + row = rng.dirichlet(np.ones(3)) + valid_value = np.broadcast_to(row, (5, 3)) + valid_logp = logp_bcast_x.eval({bcast_x_value: valid_value}) + assert valid_logp.shape == () + np.testing.assert_allclose( + valid_logp, + st.dirichlet(np.ones(3)).logpdf(row), + ) + + invalid_value = rng.dirichlet(np.ones(3), size=(5,)) + np.testing.assert_array_equal( + logp_bcast_x.eval({bcast_x_value: invalid_value}), + -np.inf, + ) + + +def test_broadcast_not_measurable_behind_other_ops(): + # The broadcast dimensions are degenerate copies; other rewrites would treat them + # as independent entries (e.g., counting the jacobian of the exp once per copy), + # so the broadcast is only measurable when directly valued + x = pt.random.normal() + y = pt.exp(pt.broadcast_to(x, (3,))) + with pytest.raises(NotImplementedError): + logp(y, y.clone()) diff --git a/tests/logprob/test_tensor.py b/tests/logprob/test_tensor.py index 1226a30e26..fc7a2bba20 100644 --- a/tests/logprob/test_tensor.py +++ b/tests/logprob/test_tensor.py @@ -48,7 +48,6 @@ from pymc.testing import assert_no_rvs -@pytest.mark.xfail(RuntimeError, reason="logprob for broadcasted RVs not implemented") def test_bcast_rv_logp(): """Test that derived logp for broadcasted RV is correct""" @@ -61,11 +60,11 @@ def test_bcast_rv_logp(): logp_combined = pt.add(*logp.values()) valid_logp = logp_combined.eval({broadcasted_x_vv: [0, 0]}) + # The broadcast dimension is consumed like a support dimension assert valid_logp.shape == () assert np.isclose(valid_logp, st.norm.logpdf(0)) # It's not possible for broadcasted dimensions to have different values - # This should either raise or return -inf invalid_logp = logp_combined.eval({broadcasted_x_vv: [0, 1]}) assert invalid_logp == -np.inf