Add Dirichlet Multinomial conjugacy (#345)

Author: Sam Anklesaria

funsor/torch/distributions.py

@@ -2,12 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
+from typing import Tuple, Union
 
 import pyro.distributions as dist
 import pyro.distributions.testing.fakes as fakes
 from pyro.distributions.torch_distribution import MaskedDistribution
 import torch
 
+from funsor.cnf import Contraction
 from funsor.distribution import (  # noqa: F401
     Bernoulli,
     FUNSOR_DIST_NAMES,
@@ -31,8 +33,9 @@
     transformeddist_to_funsor,
 )
 from funsor.domains import reals
+import funsor.ops as ops
 from funsor.tensor import Tensor, dummy_numeric_array
-from funsor.terms import Funsor, Variable, eager, to_funsor
+from funsor.terms import Binary, Funsor, Variable, eager, to_funsor
 
 
 __all__ = list(x[0] for x in FUNSOR_DIST_NAMES)
@@ -137,3 +140,32 @@ def bernoulli_to_funsor(pyro_dist, output=None, dim_to_name=None):
 eager.register(Delta, Variable, Variable, Variable)(eager_delta_variable_variable)  # noqa: F821
 eager.register(Normal, Funsor, Tensor, Funsor)(eager_normal)  # noqa: F821
 eager.register(MultivariateNormal, Funsor, Tensor, Funsor)(eager_mvn)  # noqa: F821
+
+
+@eager.register(Contraction, ops.LogAddExpOp, ops.AddOp, frozenset, Dirichlet, Multinomial)  # noqa: F821
+def eager_dirichlet_multinomial(red_op, bin_op, reduced_vars, x, y):
+    dirichlet_reduction = frozenset(x.inputs).intersection(reduced_vars)
+    if dirichlet_reduction:
+        return DirichletMultinomial(concentration=x.concentration,  # noqa: F821
+                                    total_count=y.total_count,
+                                    value=y.value)
+    else:
+        return eager(Contraction, red_op, bin_op, reduced_vars, (x, y))
+
+
+JointDirichletMultinomial = Contraction[
+    Union[ops.LogAddExpOp, ops.NullOp],
+    ops.AddOp,
+    frozenset,
+    Tuple[Dirichlet, Multinomial],  # noqa: F821
+]
+
+
+@eager.register(Binary, ops.SubOp, JointDirichletMultinomial, DirichletMultinomial)  # noqa: F821
+def eager_dirichlet_posterior(op, c, z):
+    if (z.concentration is c.terms[0].concentration) and (c.terms[1].total_count is z.total_count):
+        return Dirichlet(  # noqa: F821
+            concentration=z.concentration + c.terms[1].value,
+            value=c.terms[0].value)
+    else:
+        return None

test/test_distribution.py

@@ -268,6 +268,33 @@ def dirichlet_multinomial(concentration, total_count, value):
     assert_close(actual, expected)
 
 
+@pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
+@pytest.mark.parametrize('event_shape', [(2,), (4,), (5,)], ids=str)
+@pytest.mark.xfail(get_backend() != 'torch', reason="DirichletMultinmial is not implemented yet in NumPyro")
+def test_dirichlet_multinomial_conjugate(batch_shape, event_shape):
+    max_count = 10
+    batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
+    inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
+    full_shape = batch_shape + event_shape
+    prior = Variable("prior", reals(*full_shape))[batch_dims]
+    concentration = Tensor(ops.exp(randn(full_shape)), inputs)
+    value_data = ops.astype(randint(0, max_count, size=full_shape), 'float32')
+    value = Tensor(value_data, inputs)
+    total_count_data = value_data.sum(-1) + ops.astype(randint(0, max_count, size=batch_shape), 'float32')
+    total_count = Tensor(total_count_data, inputs)
+    p = dist.Dirichlet(concentration, value=prior)
+    p += dist.Multinomial(probs=prior, total_count=total_count)
+    marginalized = p.reduce(ops.logaddexp, set(["value"]))
+    assert isinstance(marginalized, dist.Dirichlet)
+    reduced = p.reduce(ops.logaddexp, set(["prior"]))
+    assert isinstance(reduced, dist.DirichletMultinomial)
+    assert_close(reduced.concentration, concentration)
+    assert_close(reduced.total_count, total_count)
+    result = (p - reduced)(value=value)
+    assert isinstance(result, dist.Dirichlet)
+    assert_close(result.concentration, concentration + value)
+
+
 @pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
 def test_lognormal_density(batch_shape):
     batch_dims = ('i', 'j', 'k')[:len(batch_shape)]