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