Repository: aesara-devs/aeppl (language: python)

Issue:
Make `DimShuffle`d variables measurable
An example of what should be measurable is the following:

```python
import aesara.tensor as at
from aeppl import joint_logprob

x = at.random.dirichlet([[1, 1, 1], [100, 100, 100]], name="x")
x_transpose = x.T

joint_logprob({x_transpose: x_transpose.copy()})
```

The transpose is one example of a `DimShuffle` under the hood, but there are three categories of operations that make use of `DimShuffle`s:

- shuffling dims;
- add dims;
- remove dims.

To my understand, each of these will require to be handled separately, but probably in a similar fashion.

CC @ricardoV94 @brandonwillard 


Discussion / hints:
Definitely useful 😉
This should be covered by `aesara.tensor.random.opt.local_dimshuffle_rv_lift` in a lot of cases, since it would lift the `DimShuffle` to the arguments of the `RandomVariable`.  

`local_dimshuffle_rv_lift` may not currently support that distribution, though, in which case the requisite changes might only involve an extension to that rewrite (e.g. one that covers multivariate distributions). (As @ricardoV94 mentioned below, this doesn't apply to Dirichlet distributions, due to the dimension-related constraints involved in the mapping from inputs to output/support.)

Also, I believe there's a clients-based restriction hard coded into `local_dimshuffle_rv_lift`; that will probably need to be made configurable, since the concerns it addresses are specific to maintaining the consistency of model graph evaluations (i.e. sampling from a model), and not our AePPL intermediate forms.
You cannot lift a dimshuffle through the core dimension of a distribution parameters, at least not for all distributions.


There is no way to lift `dirichlet([[1, 2, 3]]).T` because `dirichlet([[1], [2], [3]])` is not even valid.

And in the first example it would be valid but have a completely different meaning.

You could however extend the lift if the Dimshuffle only affects non-core dimensions 
> You could however extend the lift if the Dimshuffle only affects non-core dimensions

Yes, the current version of `local_dimshuffle_rv_lift` even takes this into consideration, which is why the multivariate cases are disabled (because they may need special logic or not be possible).
Now that I think about it, there may be a way to make it _always_ possible to lift these operations.  

The reason we can't lift the `DimShuffle` through `DirichletRV` is that `DirichletRV` fixes the last dimension of its inputs as the core/support dimension, so that `z.T` for `z = dirichlet(a)` results in a sample that has no possible corresponding `f(a)`.  However, if we parameterize `DirichletRV` on its core/support dimension we get `z == dirichlet(a, c=-1)` and `z.T == dirichlet(a.T, c=0)`.
Why not just DimShuffle the value variables and evaluate them at the original un-DimShuffled measurable variables?

That will also work for measurable variables that are not pure RandomVariables in general.
> Why not just DimShuffle the value variables and evaluate it at the original untransposed variable?

While that might work for the purpose of constructing a log-probability function&mdash;in the same way that `incsubtensor_rv_replace` does&mdash;it won't help with the more general issues produced by `DimShuffle`s.  

For instance, `DimShuffle`s seriously complicate the identification of subgraphs, because one needs to account for both their presence and lack thereof when matching.  A `DimShuffle` in "front" of a non-valued, measurable mixture component can prevent AePPL from identifying it as measurable, and a value-variable-based approach wouldn't help.

The problems associated with lifting `DimShuffle`s through `RandomVariable`s are also very related to the problems involved with lifting all other `Op`s through `RandomVariable`s, so, solutions to the latter can be solutions to the former, and lifting other `Op`s can have much clearer performance benefits (e.g. https://github.com/aesara-devs/aesara/issues/151).

In other words, `DimShuffle` lifting is more broadly applicable and beneficial, so almost any reason to improve those capabilities is a good one.

> That will also work for measurable variables that are not pure RandomVariables in general.

Those `RandomVariable` lifting rewrites can&mdash;and should&mdash;be refactored to apply to arbitrary measurable variables.  Currently, the rewrites only effectively make use of a `RandomVariable`'s `size` parameter, which can just as well be ignored/empty; otherwise, only the `Op`'s inputs are needed.  

When we make the multivariate extensions mentioned here, `RandomVariable.ndim_supp` and `RandomVariable.ndim_params` information will probably be needed, but those are properties we could and should be adding to our `MeasurableVariable` interface anyway.  Again, that information is really only necessary for multivariate measurable variables, so those additions wouldn't even affect most/all of our current non-`RandomVariable` `MeasurableVariable` classes.
>  A `DimShuffle` in "front" of a non-valued, measurable mixture component can prevent AePPL from identifying it as measurable, and a value-variable-based approach wouldn't help.

I don't see why this would be the case with our current eager IR. A dimshuffle of a measurable component would be converted into a MeasurableDimShuffle eventually (since such rewrite can always be applied), and then the MeasurableMixture should behave exactly as if the component was any other MeasurableVariable.

This rewrite can always be applied later if it complicates other forms of pattern matching (which I don't think is yet the case).

> The problems associated with lifting `DimShuffle`s through `RandomVariable`s are also very related to the problems involved with lifting all other `Op`s through `RandomVariable`s, so, solutions to the latter can be solutions to the former, and lifting other `Op`s can have much clearer performance benefits (e.g. https://github.com/aesara-devs/aesara/issues/151).

I am not sure about the performance benefits of lifting a (non-reducing) operator over a function of multiple inputs. For instance for the DS case you usually would need to apply it to all inputs above the RV, instead of only applying it once to the output of the RV.


> When we make the multivariate extensions mentioned here, `RandomVariable.ndim_supp` and `RandomVariable.ndim_params` information will probably be needed, but those are properties we could and should be adding to our `MeasurableVariable` interface anyway. 

This makes a lot of sense. For instance, here I had to guess what the ndim_supp of a MeasurableVariable was by how many dims were lost, but it would have been much better to have that information in the IR op level: https://github.com/aesara-devs/aeppl/blob/394dc72a536f137fbeded74a196769baba807dc9/aeppl/tensor.py#L124


> I don't see why this would be the case with our current eager IR. A dimshuffle of a measurable component would be converted into a MeasurableDimShuffle eventually (since such rewrite can always be applied), and then the MeasurableMixture should behave exactly as if the component was any other MeasurableVariable.

Yes, and then we're locked into the current eager approach.  This is the kind of coupling between design decisions that constricts our possibilities going forward.  At the same time, `MeasurableDimShuffle`s would continue to complicate rewrites on our AePPL IR, as regular `DimShuffle`s do.

> I am not sure about the performance benefits of lifting a (non-reducing) operator over a function of multiple inputs. For instance for the DS case you usually would need to apply it to all inputs above the RV, instead of only applying it once to the output of the RV.

Just like the example to which I linked, `Subtensor` lifting can reduce/remove unnecessary sampling.  That's just one simple and explicitly performance enhancing lift.  Remember, I wasn't just talking about `DimShuffle` lifts, but _any_ lifts that are facilitated by figuring out the same underlying dimension-based logic for multivariate `RandomVariable`s.  Otherwise, lifting can reduce the number of run-time transposes by bringing `Op`s closer to constant inputs, prevent copies or poor performance due to array contiguity changes, etc.

In other words, it's a type of `Op` "mobility" that we need for performance-related rewriting in general.  

If we can do things in a way that open us up to more features, approaches, _and_ solves our immediate problem(s), then we need to consider such ways the most seriously.

Fix patch:
diff --git a/aeppl/tensor.py b/aeppl/tensor.py
index 1a448a07..d8b12dcb 100644
--- a/aeppl/tensor.py
+++ b/aeppl/tensor.py
@@ -5,6 +5,7 @@
 from aesara.graph.op import compute_test_value
 from aesara.graph.opt import local_optimizer
 from aesara.tensor.basic import Join, MakeVector
+from aesara.tensor.elemwise import DimShuffle
 from aesara.tensor.extra_ops import BroadcastTo
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.opt import local_dimshuffle_rv_lift, local_rv_size_lift
@@ -98,7 +99,7 @@ def logprob_make_vector(op, values, *base_vars, **kwargs):
 
 
 class MeasurableJoin(Join):
-    """A placeholder used to specify a log-likelihood for a cumsum sub-graph."""
+    """A placeholder used to specify a log-likelihood for a join sub-graph."""
 
 
 MeasurableVariable.register(MeasurableJoin)
@@ -185,10 +186,90 @@ def find_measurable_stacks(
     return [measurable_stack]
 
 
+class MeasurableDimShuffle(DimShuffle):
+    """A placeholder used to specify a log-likelihood for a dimshuffle sub-graph."""
+
+    # Need to get the absolute path of `c_func_file`, otherwise it tries to
+    # find it locally and fails when a new `Op` is initialized
+    c_func_file = DimShuffle.get_path(DimShuffle.c_func_file)
+
+
+MeasurableVariable.register(MeasurableDimShuffle)
+
+
+@_logprob.register(MeasurableDimShuffle)
+def logprob_dimshuffle(op, values, base_var, **kwargs):
+    """Compute the log-likelihood graph for a `MeasurableDimShuffle`."""
+    (value,) = values
+
+    # Reverse the effects of dimshuffle on the value variable
+    # First, drop any augmented dimensions and reinsert any dropped dimensions
+    undo_ds: List[Union[int, str]] = [i for i, o in enumerate(op.new_order) if o != "x"]
+    dropped_dims = tuple(sorted(set(op.transposition) - set(op.shuffle)))
+    for dropped_dim in dropped_dims:
+        undo_ds.insert(dropped_dim, "x")
+    value = value.dimshuffle(undo_ds)
+
+    # Then, reshuffle remaining dims
+    undo_ds = op.shuffle
+    for dropped_dim in dropped_dims:
+        undo_ds.insert(dropped_dim, dropped_dim)
+    value = value.dimshuffle(undo_ds)
+
+    raw_logp = logprob(base_var, value)
+
+    # Re-apply original dimshuffle, ignoring any support dimensions consumed by
+    # the logprob function. This assumes that support dimensions are always in
+    # the rightmost positions, and all we need to do is to discard the highest
+    # indexes in the original dimshuffle order. Otherwise, there is no way of
+    # knowing which dimensions were consumed by the logprob function.
+    redo_ds = [o for o in op.new_order if o == "x" or o < raw_logp.ndim]
+    return raw_logp.dimshuffle(redo_ds)
+
+
+@local_optimizer([DimShuffle])
+def find_measurable_dimshuffles(fgraph, node) -> Optional[List[MeasurableDimShuffle]]:
+    r"""Finds `Dimshuffle`\s for which a `logprob` can be computed."""
+
+    if isinstance(node.op, MeasurableDimShuffle):
+        return None  # pragma: no cover
+
+    rv_map_feature: PreserveRVMappings = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    base_var = node.inputs[0]
+
+    if not (
+        base_var.owner
+        and isinstance(base_var.owner.op, MeasurableVariable)
+        and base_var not in rv_map_feature.rv_values
+    ):
+        return None  # pragma: no cover
+
+    # Make base_vars unmeasurable
+    base_var = assign_custom_measurable_outputs(base_var.owner)
+
+    measurable_dimshuffle = MeasurableDimShuffle(
+        node.op.input_broadcastable, node.op.new_order
+    )(base_var)
+    measurable_dimshuffle.name = node.outputs[0].name
+
+    return [measurable_dimshuffle]
+
+
 measurable_ir_rewrites_db.register(
     "dimshuffle_lift", local_dimshuffle_rv_lift, -5, "basic", "tensor"
 )
 
+
+# We register this later than `dimshuffle_lift` so that it is only applied as a fallback
+measurable_ir_rewrites_db.register(
+    "find_measurable_dimshuffles", find_measurable_dimshuffles, 0, "basic", "tensor"
+)
+
+
 measurable_ir_rewrites_db.register(
     "broadcast_to_lift", naive_bcast_rv_lift, -5, "basic", "tensor"
 )


Test patch:
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index 546cbe52..5f4118ba 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -1,3 +1,4 @@
+import aesara
 import numpy as np
 import pytest
 from aesara import tensor as at
@@ -188,7 +189,6 @@ def test_measurable_join_multivariate(
         y_testval = np.concatenate((base1_testval, base2_testval), axis=axis)
     else:
         y_testval = np.stack((base1_testval, base2_testval), axis=axis)
-    print(base_logps.eval({base1_vv: base1_testval, base2_vv: base2_testval}).shape)
     np.testing.assert_allclose(
         base_logps.eval({base1_vv: base1_testval, base2_vv: base2_testval}),
         y_logp.eval({y_vv: y_testval}),
@@ -205,3 +205,53 @@ def test_join_mixed_ndim_supp():
         ValueError, match="Joined logps have different number of dimensions"
     ):
         joint_logprob({y_rv: y_vv})
+
+
+@aesara.config.change_flags(cxx="")
+@pytest.mark.parametrize(
+    "ds_order",
+    [
+        (0, 2, 1),  # Swap
+        (2, 1, 0),  # Swap
+        (0, 1, 2, "x"),  # Expand
+        ("x", 0, 1, 2),  # Expand
+        (
+            0,
+            2,
+        ),  # Drop
+        (2, 0),  # Swap and drop
+        (2, 1, "x", 0),  # Swap and expand
+        ("x", 0, 2),  # Expand and drop
+        (2, "x", 0),  # Swap, expand and drop
+    ],
+)
+@pytest.mark.parametrize("multivariate", (False, True))
+def test_measurable_dimshuffle(ds_order, multivariate):
+    if multivariate:
+        base_rv = at.random.dirichlet([1, 2, 3], size=(2, 1))
+    else:
+        base_rv = at.exp(at.random.beta(1, 2, size=(2, 1, 3)))
+
+    ds_rv = base_rv.dimshuffle(ds_order)
+    base_vv = base_rv.clone()
+    ds_vv = ds_rv.clone()
+
+    # Remove support dimension axis from ds_order (i.e., 2, for multivariate)
+    if multivariate:
+        logp_ds_order = [o for o in ds_order if o == "x" or o < 2]
+    else:
+        logp_ds_order = ds_order
+
+    ref_logp = joint_logprob({base_rv: base_vv}, sum=False).dimshuffle(logp_ds_order)
+    ds_logp = joint_logprob({ds_rv: ds_vv}, sum=False)
+    assert ds_logp is not None
+
+    ref_logp_fn = aesara.function([base_vv], ref_logp)
+    ds_logp_fn = aesara.function([ds_vv], ds_logp)
+
+    base_test_value = base_rv.eval()
+    ds_test_value = at.constant(base_test_value).dimshuffle(ds_order).eval()
+
+    np.testing.assert_array_equal(
+        ref_logp_fn(base_test_value), ds_logp_fn(ds_test_value)
+    )