Support conjugacy for beta plate binomial

funsor/distribution.py

@@ -599,6 +599,24 @@ def eager_dirichlet_multinomial(red_op, bin_op, reduced_vars, x, y):
         return eager(Contraction, red_op, bin_op, reduced_vars, (x, y))
 
 
+def eager_plate_multinomial(op, x, reduced_vars):
+    if not reduced_vars.isdisjoint(x.probs.inputs):
+        return None
+    if not reduced_vars.issubset(x.value.inputs):
+        return None
+
+    backend_dist = import_module(BACKEND_TO_DISTRIBUTIONS_BACKEND[get_backend()])
+    total_count = x.total_count
+    for v in reduced_vars:
+        if v in total_count.inputs:
+            total_count = total_count.reduce(ops.add, v)
+        else:
+            total_count = total_count * x.inputs[v].size
+    return backend_dist.Multinomial(total_count=total_count,
+                                    probs=x.probs,
+                                    value=x.value.reduce(ops.add, reduced_vars))
+
+
 def _log_beta(x, y):
     return ops.lgamma(x) + ops.lgamma(y) - ops.lgamma(x + y)
 

funsor/jax/distributions.py

@@ -29,6 +29,7 @@
     eager_multinomial,
     eager_mvn,
     eager_normal,
+    eager_plate_multinomial,
     indepdist_to_funsor,
     make_dist,
     maskeddist_to_funsor,
@@ -37,7 +38,7 @@
 from funsor.domains import Real, Reals
 import funsor.ops as ops
 from funsor.tensor import Tensor, dummy_numeric_array
-from funsor.terms import Binary, Funsor, Variable, eager, to_data, to_funsor
+from funsor.terms import Binary, Funsor, Reduce, Variable, eager, to_data, to_funsor
 from funsor.util import methodof
 
 
@@ -286,6 +287,7 @@ def deltadist_to_funsor(pyro_dist, output=None, dim_to_name=None):
 if hasattr(dist, "DirichletMultinomial"):
     eager.register(Binary, ops.SubOp, JointDirichletMultinomial, DirichletMultinomial)(  # noqa: F821
         eager_dirichlet_posterior)
-
+eager.register(Reduce, ops.AddOp, Multinomial[Tensor, Funsor, Funsor], frozenset)(  # noqa: F821
+    eager_plate_multinomial)
 
 __all__ = list(x[0] for x in FUNSOR_DIST_NAMES if _get_numpyro_dist(x[0]) is not None)

funsor/memoize.py

@@ -1,7 +1,7 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections import Hashable
+from collections.abc import Hashable
 from contextlib import contextmanager
 
 import funsor.interpreter as interpreter

funsor/terms.py

@@ -7,7 +7,8 @@
 import numbers
 import typing
 import warnings
-from collections import Hashable, OrderedDict
+from collections import OrderedDict
+from collections.abc import Hashable
 from functools import reduce, singledispatch
 from weakref import WeakValueDictionary
 

funsor/torch/distributions.py

@@ -33,6 +33,7 @@
     eager_multinomial,
     eager_mvn,
     eager_normal,
+    eager_plate_multinomial,
     indepdist_to_funsor,
     make_dist,
     maskeddist_to_funsor,
@@ -41,7 +42,7 @@
 from funsor.domains import Real, Reals
 import funsor.ops as ops
 from funsor.tensor import Tensor, dummy_numeric_array
-from funsor.terms import Binary, Funsor, Variable, eager, to_data, to_funsor
+from funsor.terms import Binary, Funsor, Reduce, Variable, eager, to_data, to_funsor
 from funsor.util import methodof
 
 
@@ -275,3 +276,5 @@ def deltadist_to_funsor(pyro_dist, output=None, dim_to_name=None):
     eager_gamma_poisson)
 eager.register(Binary, ops.SubOp, JointDirichletMultinomial, DirichletMultinomial)(  # noqa: F821
     eager_dirichlet_posterior)
+eager.register(Reduce, ops.AddOp, Multinomial[Tensor, Funsor, Funsor], frozenset)(  # noqa: F821
+    eager_plate_multinomial)

test/test_distribution.py

@@ -1037,6 +1037,30 @@ def test_dirichlet_categorical_conjugate(batch_shape, size):
     _assert_conjugate_density_ok(latent, conditional, obs)
 
 
+@pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
+@pytest.mark.parametrize('size', [2, 4, 5], ids=str)
+def test_dirichlet_multinomial_conjugate_plate(batch_shape, size):
+    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 + (size,)
+    prior = Variable("prior", Reals[size])
+    concentration = Tensor(ops.exp(randn(full_shape)), inputs)
+    value_data = ops.astype(randint(0, max_count, size=batch_shape + (7, size)), 'float32')
+    obs_inputs = inputs.copy()
+    obs_inputs['plate'] = Bint[7]
+    obs = Tensor(value_data, obs_inputs)
+    total_count_data = value_data.sum(-1)
+    total_count = Tensor(total_count_data, obs_inputs)
+    latent = dist.Dirichlet(concentration, value=prior)
+    conditional = dist.Multinomial(probs=prior, total_count=total_count, value=obs)
+    p = latent + conditional.reduce(ops.add, 'plate')
+    reduced = p.reduce(ops.logaddexp, 'prior')
+    assert isinstance(reduced, Tensor)
+
+    _assert_conjugate_density_ok(latent, conditional, obs)
+
+
 @pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
 @pytest.mark.parametrize('size', [2, 4, 5], ids=str)
 def test_dirichlet_multinomial_conjugate(batch_shape, size):