Infer logprob of join_dims and split_dims

Inverted by applying the opposite operation to the value and re-applying
the original to the logp, mirroring MeasurableDimShuffle, including its
truncation of support dimensions consumed by the logp.

Claims preserve the assumption that support axes are the rightmost
positions: splits only inflate a single axis and joins contained in the
batch axes (or in the support axes, which deflate into fewer rightmost
axes) are always fine, so both are claimed unconditionally, using the
elemwise root's ndim_supp to locate the batch/support boundary. Joins
that straddle the boundary merge batch and support axes into one and are
only claimed when directly valued (#6360).

The local_join_dims/local_split_dims canonicalizations are excluded from
the logprob IR pipeline: they unconditionally lower to Reshape, which
would erase the ops before the measurable rewrites see them (and erase
the measurable subclasses after). Once pytensor canonicalizes provable
Reshapes into JoinDims/SplitDims (pymc-devs/pytensor#2280), reshape,
ravel and flatten of measurable variables will be covered by the same
rewrite.

Author: ricardoV94

pymc/logprob/rewriting.py

@@ -167,9 +167,19 @@ def remove_DiracDelta(fgraph, node):
     "basic",
     position=0.9,
 )
+# local_join_dims/local_split_dims are excluded so that JoinDims/SplitDims survive
+# until the measurable rewrites (and their measurable subclasses survive after);
+# non-measurable ones are still lowered to Reshape when the logp graph is compiled
+CANONICALIZE_IR_QUERY_ARGS = (
+    "+canonicalize",
+    "-local_eager_useless_unbatched_blockwise",
+    "-local_join_dims",
+    "-local_split_dims",
+)
+
 logprob_rewrites_db.register(
     "pre-canonicalize",
-    optdb.query("+canonicalize", "-local_eager_useless_unbatched_blockwise"),
+    optdb.query(*CANONICALIZE_IR_QUERY_ARGS),
     "basic",
     position=1,
 )
@@ -205,7 +215,7 @@ def remove_DiracDelta(fgraph, node):
 
 logprob_rewrites_db.register(
     "post-canonicalize",
-    optdb.query("+canonicalize", "-local_eager_useless_unbatched_blockwise"),
+    optdb.query(*CANONICALIZE_IR_QUERY_ARGS),
     "basic",
     position=4,
 )

pymc/logprob/tensor.py

@@ -52,6 +52,7 @@
 from pytensor.tensor.random.rewriting import (
     local_dimshuffle_rv_lift,
 )
+from pytensor.tensor.reshape import JoinDims, SplitDims, join_dims, split_dims
 from pytensor.tensor.rewriting.basic import elemwise_of
 
 from pymc.logprob.abstract import (
@@ -79,7 +80,7 @@
     get_related_valued_nodes,
     replace_rvs_by_values,
 )
-from pymc.pytensorf import constant_fold
+from pymc.pytensorf import constant_fold, get_symbolic_rv_shapes
 
 
 class MeasurableMakeVector(MeasurableOp, MakeVector):
@@ -294,24 +295,26 @@ def logprob_dimshuffle(op: MeasurableDimShuffle, values, base_var, **kwargs):
     return raw_logp.dimshuffle(redo_ds)
 
 
-def _elemwise_univariate_chain(fgraph, node) -> bool:
-    # Check whether only Elemwise operations connect a base univariate RV to the valued node through var.
+def _elemwise_root(var: TensorVariable) -> TensorVariable | None:
+    """Walk through dimension-preserving measurable operations to the root variable."""
     from pymc.distributions.distribution import SymbolicRandomVariable
     from pymc.logprob.transforms import MeasurableTransform
 
-    [inp] = node.inputs
-    [out] = node.outputs
+    if isinstance(var.owner.op, RandomVariable | SymbolicRandomVariable):
+        return var
+    elif isinstance(var.owner.op, MeasurableTransform):
+        return _elemwise_root(var.owner.inputs[var.owner.op.measurable_input_idx])
+    else:
+        return None
 
-    def elemwise_root(var: TensorVariable) -> TensorVariable | None:
-        if isinstance(var.owner.op, RandomVariable | SymbolicRandomVariable):
-            return var
-        elif isinstance(var.owner.op, MeasurableTransform):
-            return elemwise_root(var.owner.inputs[var.owner.op.measurable_input_idx])
-        else:
-            return None
+
+def _elemwise_univariate_chain(fgraph, node) -> bool:
+    # Check whether only Elemwise operations connect a base univariate RV to the valued node through var.
+    inp = node.inputs[0]
+    [out] = node.outputs
 
     # Check that the root is a univariate distribution linked by only elemwise operations
-    root = elemwise_root(inp)
+    root = _elemwise_root(inp)
     if root is None:
         return False
     elif root.owner.op.ndim_supp != 0:
@@ -590,6 +593,91 @@ def identity_icdf(op, value, base_var, **kwargs):
     return _icdf_helper(base_var, value)
 
 
+class MeasurableJoinDims(MeasurableOp, JoinDims):
+    """A placeholder used to specify a log-likelihood for a join_dims sub-graph."""
+
+
+class MeasurableSplitDims(MeasurableOp, SplitDims):
+    """A placeholder used to specify a log-likelihood for a split_dims sub-graph."""
+
+
+@node_rewriter([JoinDims, SplitDims])
+def find_measurable_join_split_dims(fgraph, node) -> list[TensorVariable] | None:
+    r"""Find `JoinDims` and `SplitDims` for which a `logprob` can be computed."""
+    if isinstance(node.op, MeasurableOp):
+        return None
+
+    base_var, *other_inputs = node.inputs
+
+    if not filter_measurable_variables([base_var]):
+        return None
+
+    if check_potential_measurability(other_inputs):
+        return None
+
+    if isinstance(node.op, JoinDims):
+        # A join that straddles the batch/support boundary merges batch and support
+        # axes into a single one, breaking the assumption that support axes are the
+        # rightmost positions, which other rewrites rely on. Such joins are only valid
+        # when directly valued. Joins contained in the batch axes (or in the support
+        # axes, which simply become fewer rightmost axes) preserve the assumption and
+        # need no restriction; splitting never straddles, since it only inflates a
+        # single axis.
+        # TODO: When we include the support axis as meta information in each intermediate
+        #  MeasurableVariable, the root walk becomes unnecessary (see https://github.com/pymc-devs/pymc/issues/6360)
+        start_axis, n_axes = node.op.start_axis, node.op.n_axes
+        root = _elemwise_root(base_var)
+        if root is None:
+            straddles = True
+        else:
+            batch_ndim = base_var.type.ndim - root.owner.op.ndim_supp
+            straddles = start_axis < batch_ndim < start_axis + n_axes
+        if straddles and not any(get_related_valued_nodes(fgraph, node)):
+            return None
+        measurable_op: MeasurableOp = MeasurableJoinDims(start_axis, n_axes)
+    else:
+        measurable_op = MeasurableSplitDims(node.op.axis)
+
+    return [measurable_op(base_var, *other_inputs)]  # type: ignore[operator]
+
+
+@_logprob.register(MeasurableJoinDims)
+def logprob_join_dims(op, values, base_var, **kwargs):
+    """Compute the log-likelihood graph for a `MeasurableJoinDims`."""
+    (value,) = values
+
+    [base_shape] = get_symbolic_rv_shapes([base_var])
+    unjoined_shape = [base_shape[i] for i in op.axis_range]
+    unjoined_value = split_dims(value, shape=unjoined_shape, axis=op.start_axis)
+
+    raw_logp = _logprob_helper(base_var, unjoined_value)
+
+    # Re-join the value dimensions, ignoring any support dimensions consumed by the
+    # logprob function (assumed to be the rightmost positions). A join lying entirely
+    # within consumed support dimensions leaves nothing to re-join, but a length-zero
+    # join (expand_dims) of remaining dimensions is still re-applied.
+    if op.start_axis > raw_logp.ndim or (op.n_axes > 0 and op.start_axis >= raw_logp.ndim):
+        return raw_logp
+    return join_dims(raw_logp, op.start_axis, min(op.n_axes, raw_logp.ndim - op.start_axis))
+
+
+@_logprob.register(MeasurableSplitDims)
+def logprob_split_dims(op, values, base_var, shape, **kwargs):
+    """Compute the log-likelihood graph for a `MeasurableSplitDims`."""
+    (value,) = values
+
+    n_axes = value.type.ndim - base_var.type.ndim + 1
+    joined_value = join_dims(value, start_axis=op.axis, n_axes=n_axes)
+
+    raw_logp = _logprob_helper(base_var, joined_value)
+
+    # Re-split the value dimensions, unless the split dimension was a support dimension
+    # consumed by the logprob function (assumed to be the rightmost positions)
+    if op.axis >= raw_logp.ndim:
+        return raw_logp
+    return split_dims(raw_logp, shape=shape, axis=op.axis)
+
+
 measurable_ir_rewrites_db.register(
     "find_measurable_casts", find_measurable_casts, "basic", "tensor"
 )
@@ -600,6 +688,10 @@ def identity_icdf(op, value, base_var, **kwargs):
     "find_measurable_identity_ops", find_measurable_identity_ops, "basic", "tensor"
 )
 
+measurable_ir_rewrites_db.register(
+    "find_measurable_join_split_dims", find_measurable_join_split_dims, "basic", "tensor"
+)
+
 
 measurable_ir_rewrites_db.register("dimshuffle_lift", local_dimshuffle_rv_lift, "basic", "tensor")
 

tests/logprob/test_tensor.py

@@ -45,6 +45,8 @@
 from pytensor.tensor.random.type import random_generator_type
 from scipy import stats as st
 
+import pymc as pm
+
 from pymc.logprob.basic import conditional_logp, icdf, logcdf, logp
 from pymc.logprob.rewriting import logprob_rewrites_db
 from pymc.testing import assert_no_rvs
@@ -765,3 +767,124 @@ def test_deep_copy(self):
             logp(y, y_vv).eval({y_vv: 0.7}),
             st.norm(0.5, 1).logpdf(0.7),
         )
+
+
+class TestMeasurableJoinSplitDims:
+    def test_join_dims(self):
+        rng = np.random.default_rng(163)
+        x = pt.random.normal(pt.arange(6).reshape((2, 3)), 1, size=(2, 3))
+        y = pt.join_dims(x)
+        y_vv = y.clone()
+
+        y_test = rng.normal(size=6)
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: y_test}),
+            st.norm(np.arange(6), 1).logpdf(y_test),
+        )
+
+    def test_split_dims(self):
+        rng = np.random.default_rng(164)
+        x = pt.random.normal(pt.arange(6), 1, size=(6,))
+        y = pt.split_dims(x, shape=(2, 3), axis=0)
+        y_vv = y.clone()
+
+        y_test = rng.normal(size=(2, 3))
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: y_test}),
+            st.norm(np.arange(6).reshape((2, 3)), 1).logpdf(y_test),
+        )
+
+    @pytest.mark.parametrize("transform_first", (False, True))
+    def test_elemwise_chain(self, transform_first):
+        rng = np.random.default_rng(165)
+        x = pt.random.normal(pt.arange(6).reshape((2, 3)), 1, size=(2, 3))
+        y = pt.join_dims(pt.exp(x)) if transform_first else pt.exp(pt.join_dims(x))
+        y_vv = y.clone()
+
+        y_test = np.abs(rng.normal(size=6)) + 0.1
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: y_test}),
+            st.norm(np.arange(6), 1).logpdf(np.log(y_test)) - np.log(y_test),
+        )
+
+    def test_multivariate_directly_valued(self):
+        rng = np.random.default_rng(166)
+
+        # The joined region extends into the support dimension consumed by the logp,
+        # so only the remaining batch dimension is re-joined
+        x = pt.random.dirichlet(pt.ones(3), size=(2,))
+        y = pt.join_dims(x)
+        y_vv = y.clone()
+        y_test = rng.dirichlet(np.ones(3), size=2).ravel()
+        y_logp = logp(y, y_vv).eval({y_vv: y_test})
+        assert y_logp.shape == (2,)
+        np.testing.assert_allclose(
+            y_logp,
+            st.dirichlet(np.ones(3)).logpdf(y_test.reshape((2, 3)).T),
+        )
+
+        # The split dimension is the support dimension itself
+        x2 = pt.random.dirichlet(pt.ones(6))
+        y2 = pt.split_dims(x2, shape=(2, 3), axis=0)
+        y2_vv = y2.clone()
+        y2_test = rng.dirichlet(np.ones(6)).reshape((2, 3))
+        y2_logp = logp(y2, y2_vv).eval({y2_vv: y2_test})
+        assert y2_logp.shape == ()
+        np.testing.assert_allclose(
+            y2_logp,
+            st.dirichlet(np.ones(6)).logpdf(y2_test.ravel()),
+        )
+
+    def test_multivariate_indirect_join_within_batch(self):
+        # A join contained in the batch axes leaves the support axes rightmost,
+        # so it is measurable even behind other operations
+        rng = np.random.default_rng(168)
+        x = pt.random.dirichlet(pt.ones(3), size=(2, 2))
+        y = pt.exp(pt.join_dims(x, start_axis=0, n_axes=2))
+        y_vv = y.clone()
+        y_test = np.exp(rng.dirichlet(np.ones(3), size=(2, 2)).reshape((4, 3)))
+        y_logp = logp(y, y_vv).eval({y_vv: y_test})
+        assert y_logp.shape == (4,)
+        np.testing.assert_allclose(
+            y_logp,
+            st.dirichlet(np.ones(3)).logpdf(np.log(y_test).T) - np.log(y_test).sum(-1),
+        )
+
+    def test_multivariate_indirect_join_within_support(self):
+        # A join contained in the support axes just deflates them into fewer
+        # rightmost axes, so it is measurable even behind other operations
+        rng = np.random.default_rng(169)
+        x = pm.MatrixNormal.dist(mu=np.zeros((2, 3)), rowcov=np.eye(2), colcov=np.eye(3))
+        y = pt.exp(pt.join_dims(x, start_axis=0, n_axes=2))
+        y_vv = y.clone()
+        y_test = np.exp(rng.normal(size=6))
+        # With identity covariances the matrix normal entries are iid standard normal
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: y_test}),
+            st.norm.logpdf(np.log(y_test)).sum() - np.log(y_test).sum(),
+        )
+
+    # batch size 1 is the treacherous case: if the straddling join were wrongly
+    # claimed, the truncated logp would silently broadcast against the fused value
+    # dimension instead of raising a shape error
+    @pytest.mark.parametrize("batch_size", (1, 2))
+    def test_multivariate_indirect_straddling_join_not_measurable(self, batch_size):
+        # A join straddling the batch/support boundary merges batch and support axes;
+        # it is only measurable when directly valued
+        x = pt.random.dirichlet(pt.ones(3), size=(batch_size,))
+        y = pt.exp(pt.join_dims(x))
+        with pytest.raises(NotImplementedError):
+            logp(y, y.clone())
+
+    def test_multivariate_indirect_split(self):
+        # Splitting only inflates a single axis, so it is measurable even for
+        # multivariate variables behind other operations
+        rng = np.random.default_rng(167)
+        x = pt.random.dirichlet(pt.ones(6))
+        y = pt.exp(pt.split_dims(x, shape=(2, 3), axis=0))
+        y_vv = y.clone()
+        y_test = np.exp(rng.dirichlet(np.ones(6)).reshape((2, 3)))
+        np.testing.assert_allclose(
+            logp(y, y_vv).eval({y_vv: y_test}),
+            st.dirichlet(np.ones(6)).logpdf(np.log(y_test).ravel()) - np.log(y_test).sum(),
+        )