Repository: pymc-devs/pymc
Pull request: Default `moment` for `CustomDist` provided with a `dist` function
Add default implementation of `moment` for `CustomDist` with a `dist` function and close #6804


<!-- readthedocs-preview pymc start -->
----
:books: Documentation preview :books:: https://pymc--6873.org.readthedocs.build/en/6873/

<!-- readthedocs-preview pymc end -->

Problem statement:
Implement default `moment` function for `CustomDist` provided with a `dist` function
### Description

```python
import pymc as pm

def dist(mu, sigma, size):
  return pm.math.exp(pm.Normal.dist(mu, sigma size=size))`

with pm.Model() as m:
  ...
  x = pm.CustomDist("x", mu, sigma, dist=dist)

m.initial_point()
```

Since the `dist` is a symbolic forward graph, we could implement the `moment` automatically as
```python
def moment_dist(mu, sigma, size):
  return pm.math.exp(moment(pm.Normal.dist(mu, sigma, size))
```

Would just need to replace the RVs in the `dist` graph, by their moments.

Hints from issue discussion:
I'm kind of late to this issue but I need to comment that the expected value of the exponential of a Normally distributed random variable is not the exponential of the normal's location.

```python
>>> moment(pm.math.exp(pm.Normal.dist(mu, sigma, size)))
pm.math.exp(mu + sigma **2 / 2)
```

Just like the `Lognormal` distribution's moment. I do like the idea of being able to deduce the moments from the data generation process, and I like the approach of using a rewrite to do so, as what we do with `logp` graphs. The problem is that moments are integrals on the altered measures of some basic random variables. Maybe @maresb has some input on this topic, but to me it sounds like a mathematically challenging endeavor.
With moment we're just looking for a nonzero density point to start sampling. We can call it something else, but definitely not confusing it with mean.

Commit message: add comments

Patch to review:
diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
index 6bf39ed998..4296c6306a 100644
--- a/pymc/distributions/distribution.py
+++ b/pymc/distributions/distribution.py
@@ -23,16 +23,19 @@
 
 import numpy as np
 
+from pytensor import shared
 from pytensor import tensor as pt
 from pytensor.compile.builders import OpFromGraph
-from pytensor.graph import FunctionGraph, node_rewriter
-from pytensor.graph.basic import Node, Variable
-from pytensor.graph.replace import clone_replace
-from pytensor.graph.rewriting.basic import in2out
+from pytensor.graph import FunctionGraph, clone_replace, node_rewriter
+from pytensor.graph.basic import Node, Variable, io_toposort
+from pytensor.graph.features import ReplaceValidate
+from pytensor.graph.rewriting.basic import GraphRewriter, in2out
 from pytensor.graph.utils import MetaType
+from pytensor.scan.op import Scan
 from pytensor.tensor.basic import as_tensor_variable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.rewriting import local_subtensor_rv_lift
+from pytensor.tensor.random.type import RandomGeneratorType, RandomType
 from pytensor.tensor.random.utils import normalize_size_param
 from pytensor.tensor.rewriting.shape import ShapeFeature
 from pytensor.tensor.variable import TensorVariable
@@ -83,6 +86,64 @@
 PLATFORM = sys.platform
 
 
+def filter_RNGs(params):
+    return [p for p in params if not isinstance(p.type, (RandomType, RandomGeneratorType))]
+
+
+class MomentRewrite(GraphRewriter):
+    def rewrite_moment_scan_node(self, node):
+        if not isinstance(node.op, Scan):
+            return
+
+        node_inputs, node_outputs = node.op.inner_inputs, node.op.inner_outputs
+        op = node.op
+
+        local_fgraph_topo = io_toposort(node_inputs, node_outputs)
+
+        replace_with_moment = []
+        to_replace_set = set()
+
+        for nd in local_fgraph_topo:
+            if nd not in to_replace_set and isinstance(
+                nd.op, (RandomVariable, SymbolicRandomVariable)
+            ):
+                replace_with_moment.append(nd.out)
+                to_replace_set.add(nd)
+        givens = {}
+        if len(replace_with_moment) > 0:
+            for item in replace_with_moment:
+                givens[item] = moment(item)
+        else:
+            return
+        op_outs = clone_replace(node_outputs, replace=givens)
+
+        nwScan = Scan(
+            node_inputs,
+            op_outs,
+            op.info,
+            mode=op.mode,
+            profile=op.profile,
+            truncate_gradient=op.truncate_gradient,
+            name=op.name,
+            allow_gc=op.allow_gc,
+        )
+        nw_node = nwScan(*(node.inputs), return_list=True)[0].owner
+        return nw_node
+
+    def add_requirements(self, fgraph):
+        fgraph.attach_feature(ReplaceValidate())
+
+    def apply(self, fgraph):
+        for node in fgraph.toposort():
+            if isinstance(node.op, (RandomVariable, SymbolicRandomVariable)):
+                fgraph.replace(node.out, moment(node.out))
+            elif isinstance(node.op, Scan):
+                new_node = self.rewrite_moment_scan_node(node)
+                if new_node is not None:
+                    for out1, out2 in zip(node.outputs, new_node.outputs):
+                        fgraph.replace(out1, out2)
+
+
 class _Unpickling:
     pass
 
@@ -622,13 +683,28 @@ def dist(
         if logcdf is None:
             logcdf = default_not_implemented(class_name, "logcdf")
 
+        def dist_moment(rv, size, *dist_params):
+            fgraph = rv.owner.op.fgraph.clone()
+            replace_moments = MomentRewrite()
+            replace_moments.rewrite(fgraph)
+            # we need to replace dummy variable by actual dist params
+            for i, par in enumerate([size] + list(dist_params)):
+                fgraph.replace(fgraph.inputs[i], par)
+            # we need to replace dymmy random generators in Scan node inputs
+            for node in fgraph.toposort():
+                if isinstance(node.op, Scan):
+                    for inp in node.inputs:
+                        if isinstance(inp.type, RandomGeneratorType):
+                            fgraph.replace(
+                                inp,
+                                shared(np.random.Generator(np.random.PCG64())),
+                                import_missing=True,
+                            )
+            moment = fgraph.outputs[-1]
+            return moment
+
         if moment is None:
-            moment = functools.partial(
-                default_moment,
-                rv_name=class_name,
-                has_fallback=True,
-                ndim_supp=ndim_supp,
-            )
+            moment = dist_moment
 
         return super().dist(
             dist_params,
@@ -687,7 +763,7 @@ def custom_dist_logcdf(op, value, size, *params, **kwargs):
 
         @_moment.register(rv_type)
         def custom_dist_get_moment(op, rv, size, *params):
-            return moment(rv, size, *params[: len(params)])
+            return moment(rv, size, *filter_RNGs(params))
 
         @_change_dist_size.register(rv_type)
         def change_custom_symbolic_dist_size(op, rv, new_size, expand):
diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py
index 9bdfdb2053..8c16559581 100644
--- a/tests/distributions/test_distribution.py
+++ b/tests/distributions/test_distribution.py
@@ -430,6 +430,85 @@ def custom_dist(mu, sigma, size):
         ip = m.initial_point()
         np.testing.assert_allclose(m.compile_logp()(ip), ref_m.compile_logp()(ip))
 
+    @pytest.mark.parametrize(
+        "dist_params, size, expected, dist_fn",
+        [
+            (
+                (5, 1),
+                None,
+                np.exp(5),
+                lambda mu, sigma, size: pt.exp(pm.Normal.dist(mu, sigma, size=size)),
+            ),
+            (
+                (2, np.ones(5)),
+                None,
+                np.exp([2, 2, 2, 2, 2] + np.ones(5)),
+                lambda mu, sigma, size: pt.exp(
+                    pm.Normal.dist(mu, sigma, size=size) + pt.ones(size)
+                ),
+            ),
+            (
+                (1, 2),
+                None,
+                np.sqrt(np.exp(1 + 0.5 * 2**2)),
+                lambda mu, sigma, size: pt.sqrt(pm.LogNormal.dist(mu, sigma, size=size)),
+            ),
+            (
+                (4,),
+                (3,),
+                np.log([4, 4, 4]),
+                lambda nu, size: pt.log(pm.ChiSquared.dist(nu, size=size)),
+            ),
+            (
+                (12, 1),
+                None,
+                12,
+                lambda mu1, sigma, size: pm.Normal.dist(mu1, sigma, size=size),
+            ),
+        ],
+    )
+    def test_custom_dist_default_moment(self, dist_params, size, expected, dist_fn):
+        with Model() as model:
+            CustomDist("x", *dist_params, dist=dist_fn, size=size)
+        assert_moment_is_expected(model, expected)
+
+    def test_custom_dist_default_moment_inner_graph(self):
+        def scan_step(left, right):
+            x = pm.Uniform.dist(left, right)
+            x_update = collect_default_updates([x])
+            return x, x_update
+
+        def dist(size):
+            xs, updates = scan(
+                fn=scan_step,
+                sequences=[
+                    pt.as_tensor_variable(np.array([-4, -3])),
+                    pt.as_tensor_variable(np.array([-2, -1])),
+                ],
+                name="xs",
+            )
+            return xs
+
+        with Model() as model:
+            CustomDist("x", dist=dist)
+        assert_moment_is_expected(model, np.array([-3, -2]))
+
+    @pytest.mark.parametrize(
+        "left, right, size, expected",
+        [
+            (-1, 1, None, 0 + 5),
+            (-3, -1, None, -2 + 5),
+            (-3, 1, (3,), np.array([-1 + 5, -1 + 5, -1 + 5])),
+        ],
+    )
+    def test_custom_dist_default_moment_nested(self, left, right, size, expected):
+        def dist_fn(left, right, size):
+            return pm.Truncated.dist(pm.Normal.dist(0, 1), left, right, size=size) + 5
+
+        with Model() as model:
+            CustomDist("x", left, right, size=size, dist=dist_fn)
+        assert_moment_is_expected(model, expected)
+
     def test_logcdf_inference(self):
         def custom_dist(mu, sigma, size):
             return pt.exp(pm.Normal.dist(mu, sigma, size=size))


Code review comments:
diff_hunk: @@ -430,6 +430,85 @@ def custom_dist(mu, sigma, size):
         ip = m.initial_point()
         np.testing.assert_allclose(m.compile_logp()(ip), ref_m.compile_logp()(ip))
 
+    @pytest.mark.parametrize(
+        "dist_params, size, expected, dist_fn",
+        [
+            (
+                (5, 1),
+                None,
+                np.exp(5),
+                lambda mu, sigma, size: pt.exp(pm.Normal.dist(mu, sigma, size=size)),
+            ),
+            (
+                (2, np.ones(5)),
+                None,
+                np.exp([2, 2, 2, 2, 2] + np.ones(5)),
+                lambda mu, sigma, size: pt.exp(
+                    pm.Normal.dist(mu, sigma, size=size) + pt.ones(size)
+                ),
+            ),
+            (
+                (1, 2),
+                None,
+                np.sqrt(np.exp(1 + 0.5 * 2**2)),
+                lambda mu, sigma, size: pt.sqrt(pm.LogNormal.dist(mu, sigma, size=size)),
+            ),
+            (
+                (4,),
+                (3,),
+                np.log([4, 4, 4]),
+                lambda nu, size: pt.log(pm.ChiSquared.dist(nu, size=size)),
+            ),
+            (
+                (12, 1),
+                None,
+                12,
+                lambda mu1, sigma, size: pm.Normal.dist(mu1, sigma, size=size),
+            ),
+        ],
+    )
+    def test_custom_dist_default_moment(self, dist_params, size, expected, dist_fn):
+        with Model() as model:
+            CustomDist("x", *dist_params, dist=dist_fn, size=size)
+        assert_moment_is_expected(model, expected)
+
+    def test_custom_dist_default_moment_inner_graph(self):
path: tests/distributions/test_distribution.py
text: @user1:
```suggestion
    def test_custom_dist_default_moment_scan(self):
```

diff_hunk: @@ -622,13 +683,28 @@ def dist(
         if logcdf is None:
             logcdf = default_not_implemented(class_name, "logcdf")
 
+        def dist_moment(rv, size, *dist_params):
+            fgraph = rv.owner.op.fgraph.clone()
path: pymc/distributions/distribution.py
text: @user1:
This method can now be used for any OpFromGraph not just CustomDist ones, so I would move it to the MomentRewriter

diff_hunk: @@ -622,13 +683,28 @@ def dist(
         if logcdf is None:
             logcdf = default_not_implemented(class_name, "logcdf")
 
+        def dist_moment(rv, size, *dist_params):
+            fgraph = rv.owner.op.fgraph.clone()
+            replace_moments = MomentRewrite()
+            replace_moments.rewrite(fgraph)
+            # we need to replace dummy variable by actual dist params
+            for i, par in enumerate([size] + list(dist_params)):
+                fgraph.replace(fgraph.inputs[i], par)
+            # we need to replace dymmy random generators in Scan node inputs
+            for node in fgraph.toposort():
+                if isinstance(node.op, Scan):
path: pymc/distributions/distribution.py
text: @user1:
Why do we need a special logic for Scan inside an OpFromGraph?

@author:
@user1 After cloning fgraph we have dummy placeholders in `Scan` inner graph inputs for `RandomGenerator`, and we need to replace it by real random generator to avoid error later when we will try to calculate moment. Do we have another way to do it?

@user1:
I see, so the problem is that the scan moment rewrite does not remove the useless updates. That's probably okay. The issue is that we filtered out the RNGs!!! If we had them the replacement would have worked fine. I was able to simplify the function to this and it passes the test that has a scan:

```python
        def dist_moment(rv, *args):
            node = rv.owner
            op = node.op
            rv_out_idx = node.outputs.index(rv)
            
            fgraph = op.fgraph.clone()
            replace_moments = MomentRewrite()
            replace_moments.rewrite(fgraph)
            # Replace dummy inner inputs by outer inputs
            fgraph.replace_all(tuple(zip(op.inner_inputs, node.inputs)), import_missing=True)
            moment = fgraph.outputs[rv_out_idx]
            return moment
```

Note that I basically got a handle of the filtered RNGs by doing node.inputs. We don't need to do this if we don't filter them out.

To tidy up the code, we should dispatch this moment function on `CustomSymbolicDistRV` outside of the `_CustomSymbolicDist` class. Then we only dispatch a user provided moment when that's given like we do now with `if logp is not None:`

You can remove the moment related stuff that now exists in `_CustomSymbolicDist.dist`



@author:
@user1 Great, thank you! Will do

diff_hunk: @@ -622,13 +683,28 @@ def dist(
         if logcdf is None:
             logcdf = default_not_implemented(class_name, "logcdf")
 
+        def dist_moment(rv, size, *dist_params):
+            fgraph = rv.owner.op.fgraph.clone()
+            replace_moments = MomentRewrite()
+            replace_moments.rewrite(fgraph)
+            # we need to replace dummy variable by actual dist params
+            for i, par in enumerate([size] + list(dist_params)):
+                fgraph.replace(fgraph.inputs[i], par)
path: pymc/distributions/distribution.py
text: @user1:
You can use `fgraph.replace_all` to avoid the loop