Automate distribution testing

Makefile

@@ -43,7 +43,9 @@ ifeq (${FUNSOR_BACKEND}, torch)
 	python examples/sensor.py --seed=0 --num-frames=2 -n 1
 	@echo PASS
 else ifeq (${FUNSOR_BACKEND}, jax)
-	pytest -v -n auto --ignore=test/examples --ignore=test/pyro --ignore=test/pyroapi
+	pytest -v -n auto --ignore=test/examples --ignore=test/pyro --ignore=test/pyroapi --ignore=test/test_distribution.py --ignore=test/test_distribution_generic.py
+	pytest -v -n auto test/test_distribution.py
+	pytest -v -n auto test/test_distribution_generic.py
 	@echo PASS
 else
 	# default backend

funsor/distribution.py

@@ -134,23 +134,28 @@ def unscaled_sample(self, sampled_vars, sample_inputs, rng_key=None):
         value = params.pop("value")
         assert all(isinstance(v, (Number, Tensor)) for v in params.values())
         assert isinstance(value, Variable) and value.name in sampled_vars
-        inputs_, tensors = align_tensors(*params.values())
-        inputs = OrderedDict(sample_inputs.items())
-        inputs.update(inputs_)
-        sample_shape = tuple(v.size for v in sample_inputs.values())
 
-        raw_dist = self.dist_class(**dict(zip(self._ast_fields[:-1], tensors)))
+        value_name = value.name
+        raw_dist, value_output, dim_to_name = self._get_raw_dist()
+        for d, name in zip(range(len(sample_inputs), 0, -1), sample_inputs.keys()):
+            dim_to_name[-d - len(raw_dist.batch_shape)] = name
+
+        sample_shape = tuple(v.size for v in sample_inputs.values())
         sample_args = (sample_shape,) if get_backend() == "torch" else (rng_key, sample_shape)
         if self.has_rsample:
-            raw_sample = raw_dist.rsample(*sample_args)
+            raw_value = raw_dist.rsample(*sample_args)
         else:
-            raw_sample = ops.detach(raw_dist.sample(*sample_args))
+            raw_value = ops.detach(raw_dist.sample(*sample_args))
 
-        result = funsor.delta.Delta(value.name, Tensor(raw_sample, inputs, value.output.dtype))
+        funsor_value = to_funsor(raw_value, output=value_output, dim_to_name=dim_to_name)
+        funsor_value = funsor_value.align(
+            tuple(sample_inputs) + tuple(inp for inp in self.inputs if inp in funsor_value.inputs))
+        result = funsor.delta.Delta(value_name, funsor_value)
         if not self.has_rsample:
             # scaling of dice_factor by num samples should already be handled by Funsor.sample
-            raw_log_prob = raw_dist.log_prob(raw_sample)
-            dice_factor = Tensor(raw_log_prob - ops.detach(raw_log_prob), inputs)
+            raw_log_prob = raw_dist.log_prob(raw_value)
+            dice_factor = to_funsor(raw_log_prob - ops.detach(raw_log_prob),
+                                    output=self.output, dim_to_name=dim_to_name)
             result = result + dice_factor
         return result
 

funsor/jax/distributions.py

@@ -175,14 +175,11 @@ def _infer_param_domain(cls, name, raw_shape):
 @to_funsor.register(dist.BinomialProbs)
 @to_funsor.register(dist.BinomialLogits)
 def categorical_to_funsor(numpyro_dist, output=None, dim_to_name=None):
-    new_pyro_dist = _NumPyroWrapper_Binomial(probs=numpyro_dist.probs)
+    new_pyro_dist = _NumPyroWrapper_Binomial(total_count=numpyro_dist.total_count, probs=numpyro_dist.probs)
     return backenddist_to_funsor(Binomial, new_pyro_dist, output, dim_to_name)  # noqa: F821
 
 
 @to_funsor.register(dist.CategoricalProbs)
-# XXX: in Pyro backend, we always convert pyro.distributions.Categorical
-# to funsor.torch.distributions.Categorical
-@to_funsor.register(dist.CategoricalLogits)
 def categorical_to_funsor(numpyro_dist, output=None, dim_to_name=None):
     new_pyro_dist = _NumPyroWrapper_Categorical(probs=numpyro_dist.probs)
     return backenddist_to_funsor(Categorical, new_pyro_dist, output, dim_to_name)  # noqa: F821
@@ -191,7 +188,7 @@ def categorical_to_funsor(numpyro_dist, output=None, dim_to_name=None):
 @to_funsor.register(dist.MultinomialProbs)
 @to_funsor.register(dist.MultinomialLogits)
 def categorical_to_funsor(numpyro_dist, output=None, dim_to_name=None):
-    new_pyro_dist = _NumPyroWrapper_Multinomial(probs=numpyro_dist.probs)
+    new_pyro_dist = _NumPyroWrapper_Multinomial(total_count=numpyro_dist.total_count, probs=numpyro_dist.probs)
     return backenddist_to_funsor(Multinomial, new_pyro_dist, output, dim_to_name)  # noqa: F821
 
 

funsor/jax/ops.py

@@ -121,6 +121,11 @@ def _is_numeric_array(x):
     return True
 
 
+@ops.isnan.register(array)
+def _isnan(x):
+    return np.isnan(x)
+
+
 @ops.lgamma.register(array)
 def _lgamma(x):
     return gammaln(x)

funsor/ops/array.py

@@ -24,6 +24,7 @@
 diagonal = Op("diagonal")
 einsum = Op("einsum")
 full_like = Op(np.full_like)
+isnan = Op(np.isnan)
 prod = Op(np.prod)
 stack = Op("stack")
 sum = Op(np.sum)
@@ -300,6 +301,7 @@ def unsqueeze(x, dim):
     'finfo',
     'full_like',
     'is_numeric_array',
+    'isnan',
     'logaddexp',
     'logsumexp',
     'new_arange',

funsor/testing.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import contextlib
+import importlib
 import itertools
 import numbers
 import operator
@@ -265,6 +266,24 @@ def randn(*args):
         return np.array(np.random.randn(*shape))
 
 
+def random_scale_tril(*args):
+    if isinstance(args[0], tuple):
+        assert len(args) == 1
+        shape = args[0]
+    else:
+        shape = args
+
+    from funsor.distribution import BACKEND_TO_DISTRIBUTIONS_BACKEND
+    backend_dist = importlib.import_module(BACKEND_TO_DISTRIBUTIONS_BACKEND[get_backend()]).dist
+
+    if get_backend() == "torch":
+        data = randn(shape)
+        return backend_dist.transforms.transform_to(backend_dist.constraints.lower_cholesky)(data)
+    else:
+        data = randn(shape[:-2] + (shape[-1] * (shape[-1] + 1) // 2,))
+        return backend_dist.biject_to(backend_dist.constraints.lower_cholesky)(data)
+
+
 def zeros(*args):
     if isinstance(args[0], tuple):
         assert len(args) == 1

funsor/torch/distributions.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import functools
+import numbers
 from typing import Tuple, Union
 
 import pyro.distributions as dist
@@ -40,7 +41,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_funsor
+from funsor.terms import Binary, Funsor, Variable, eager, to_data, to_funsor
 from funsor.util import methodof
 
 
@@ -153,6 +154,19 @@ def _infer_param_domain(cls, name, raw_shape):
     return Real
 
 
+###########################################################
+# Converting distribution funsors to PyTorch distributions
+###########################################################
+
+@to_data.register(Multinomial)  # noqa: F821
+def multinomial_to_data(funsor_dist, name_to_dim=None):
+    probs = to_data(funsor_dist.probs, name_to_dim)
+    total_count = to_data(funsor_dist.total_count, name_to_dim)
+    if isinstance(total_count, numbers.Number) or len(total_count.shape) == 0:
+        return dist.Multinomial(int(total_count), probs=probs)
+    raise NotImplementedError("inhomogeneous total_count not supported")
+
+
 ###############################################
 # Converting PyTorch Distributions to funsors
 ###############################################

funsor/torch/ops.py

@@ -107,6 +107,11 @@ def _is_numeric_array(x):
     return True
 
 
+@ops.isnan.register(torch.Tensor)
+def _isnan(x):
+    return torch.isnan(x)
+
+
 @ops.lgamma.register(torch.Tensor)
 def _lgamma(x):
     return x.lgamma()

test/test_distribution.py

@@ -19,7 +19,8 @@
 from funsor.interpreter import interpretation, reinterpret
 from funsor.tensor import Einsum, Tensor, numeric_array, stack
 from funsor.terms import Independent, Variable, eager, lazy, to_funsor
-from funsor.testing import assert_close, check_funsor, rand, randint, randn, random_mvn, random_tensor, xfail_param
+from funsor.testing import assert_close, check_funsor, rand, randint, randn, \
+    random_mvn, random_scale_tril, random_tensor, xfail_param
 from funsor.util import get_backend
 
 pytestmark = pytest.mark.skipif(get_backend() == "numpy",
@@ -472,15 +473,6 @@ def test_mvn_defaults():
     assert dist.MultivariateNormal(loc, scale_tril) is dist.MultivariateNormal(loc, scale_tril, value)
 
 
-def _random_scale_tril(shape):
-    if get_backend() == "torch":
-        data = randn(shape)
-        return backend_dist.transforms.transform_to(backend_dist.constraints.lower_cholesky)(data)
-    else:
-        data = randn(shape[:-2] + (shape[-1] * (shape[-1] + 1) // 2,))
-        return backend_dist.biject_to(backend_dist.constraints.lower_cholesky)(data)
-
-
 @pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
 def test_mvn_density(batch_shape):
     batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
@@ -493,7 +485,7 @@ def mvn(loc: Reals[3], scale_tril: Reals[3, 3], value: Reals[3]) -> Real:
     check_funsor(mvn, {'loc': Reals[3], 'scale_tril': Reals[3, 3], 'value': Reals[3]}, Real)
 
     loc = Tensor(randn(batch_shape + (3,)), inputs)
-    scale_tril = Tensor(_random_scale_tril(batch_shape + (3, 3)), inputs)
+    scale_tril = Tensor(random_scale_tril(batch_shape + (3, 3)), inputs)
     value = Tensor(randn(batch_shape + (3,)), inputs)
     expected = mvn(loc, scale_tril, value)
     check_funsor(expected, inputs, Real)
@@ -509,7 +501,7 @@ def test_mvn_gaussian(batch_shape):
     inputs = OrderedDict((k, Bint[v]) for k, v in zip(batch_dims, batch_shape))
 
     loc = Tensor(randn(batch_shape + (3,)), inputs)
-    scale_tril = Tensor(_random_scale_tril(batch_shape + (3, 3)), inputs)
+    scale_tril = Tensor(random_scale_tril(batch_shape + (3, 3)), inputs)
     value = Tensor(randn(batch_shape + (3,)), inputs)
 
     expected = dist.MultivariateNormal(loc, scale_tril, value)
@@ -808,7 +800,7 @@ def test_mvn_sample(with_lazy, batch_shape, sample_inputs, event_shape):
     inputs = OrderedDict((k, Bint[v]) for k, v in zip(batch_dims, batch_shape))
 
     loc = randn(batch_shape + event_shape)
-    scale_tril = _random_scale_tril(batch_shape + event_shape * 2)
+    scale_tril = random_scale_tril(batch_shape + event_shape * 2)
     funsor_dist_class = dist.MultivariateNormal
     params = (loc, scale_tril)
 
@@ -893,7 +885,7 @@ def test_binomial_sample(with_lazy, batch_shape, sample_inputs):
     funsor_dist_class = dist.Binomial
     params = (total_count, probs)
 
-    _check_sample(funsor_dist_class, params, sample_inputs, inputs, atol=2e-2, skip_grad=True, with_lazy=with_lazy)
+    _check_sample(funsor_dist_class, params, sample_inputs, inputs, atol=5e-2, skip_grad=True, with_lazy=with_lazy)
 
 
 @pytest.mark.parametrize('sample_inputs', [(), ('ii',), ('ii', 'jj'), ('ii', 'jj', 'kk')])