Added distribution test values (pymc-devs#224)

* Adding meta executor template

* Added test values for everything except wishart

* Add wishart test_value

* Remove backend arithmetic test

* Restructured tests and added multidim parameters

* Working test_value

* Fix mypy problems and remove the meta executor skeleton that will go in a different PR

pymc4/distributions/continuous.py

@@ -1,6 +1,7 @@
 """PyMC4 continuous random variables for tensorflow."""
 import math
 
+import tensorflow as tf
 from tensorflow_probability import distributions as tfd
 from tensorflow_probability import bijectors as bij
 from pymc4.distributions.distribution import (
@@ -973,6 +974,14 @@ def upper_limit(self):
     def lower_limit(self):
         return self.conditions["scale"]
 
+    @property
+    def test_value(self):
+        return (
+            tf.zeros(self.batch_shape + self.event_shape, dtype=self.dtype)
+            + self.conditions["scale"]
+            + 1
+        )
+
 
 class StudentT(ContinuousDistribution):
     r"""Student's T random variable.

pymc4/distributions/discrete.py

@@ -295,8 +295,8 @@ def __init__(self, name, probs, **kwargs):
 
     @staticmethod
     def _init_distribution(conditions):
-        probs = conditions["probs"]
-        outcomes = tf.range(float(len(probs)))
+        probs = tf.convert_to_tensor(conditions["probs"])
+        outcomes = tf.range(probs.shape[-1])
         return tfd.FiniteDiscrete(outcomes, probs=probs)
 
     def lower_limit(self):
@@ -342,6 +342,8 @@ class Geometric(BoundedDiscreteDistribution):
     probs : float
         Probability of success on an individual trial (0 < probs <= 1).
     """
+    # Another example for a wrong type used on the tensorflow side
+    _test_value = 2.0  # type: ignore
 
     def __init__(self, name, probs, **kwargs):
         super().__init__(name, probs=probs, **kwargs)
@@ -410,6 +412,8 @@ def NegBinom(a, m, x):
     probs : float
         Probability of success on an individual trial (0 < probs <= 1).
     """
+    # For some ridiculous reason, tfp needs negative binomial values to be floats...
+    _test_value = 0.0  # type: ignore
 
     def __init__(self, name, total_count, probs, **kwargs):
         super().__init__(name, total_count=total_count, probs=probs, **kwargs)
@@ -462,6 +466,8 @@ class Poisson(PositiveDiscreteDistribution):
     The Poisson distribution can be derived as a limiting case of the
     binomial distribution.
     """
+    # For some ridiculous reason, tfp needs poisson values to be floats...
+    _test_value = 0.0  # type: ignore
 
     def __init__(self, name, rate, **kwargs):
         super().__init__(name, rate=rate, **kwargs)

pymc4/distributions/distribution.py

@@ -2,6 +2,7 @@
 import copy
 from typing import Optional, Union, Any
 
+import tensorflow as tf
 from tensorflow_probability import distributions as tfd
 from pymc4.coroutine_model import Model, unpack
 from pymc4.distributions.batchstack import BatchStacker
@@ -27,6 +28,8 @@
 class Distribution(Model):
     """Statistical distribution."""
 
+    _test_value = 0.0
+
     def __init__(
         self,
         name: Optional[NameType],
@@ -87,6 +90,10 @@ def unpack_conditions(**kwargs) -> dict:
         """
         return kwargs
 
+    @property
+    def test_value(self):
+        return tf.broadcast_to(self._test_value, self.batch_shape + self.event_shape)
+
     def sample(self, sample_shape=(), seed=None):
         """
         Forward sampling implementation.
@@ -212,11 +219,11 @@ def is_anonymous(self):
 
 
 class ContinuousDistribution(Distribution):
-    ...
+    _test_value = 0.0
 
 
 class DiscreteDistribution(Distribution):
-    ...
+    _test_value = 0
 
 
 class BoundedDistribution(Distribution):
@@ -230,11 +237,15 @@ def upper_limit(self):
 
 
 class BoundedDiscreteDistribution(DiscreteDistribution, BoundedDistribution):
-    ...
+    @property
+    def _test_value(self):
+        return tf.cast(tf.round(0.5 * (self.upper_limit() + self.lower_limit())), self.dtype)
 
 
 class BoundedContinuousDistribution(ContinuousDistribution, BoundedDistribution):
-    ...
+    @property
+    def _test_value(self):
+        return 0.5 * (self.upper_limit() + self.lower_limit())
 
 
 class UnitContinuousDistribution(BoundedContinuousDistribution):
@@ -252,6 +263,8 @@ def upper_limit(self):
 
 
 class PositiveContinuousDistribution(BoundedContinuousDistribution):
+    _test_value = 1.0
+
     def _init_transform(self, transform):
         if transform is None:
             return transforms.Log()
@@ -266,6 +279,8 @@ def upper_limit(self):
 
 
 class PositiveDiscreteDistribution(BoundedDiscreteDistribution):
+    _test_value = 1
+
     def lower_limit(self):
         return 0
 
@@ -274,4 +289,6 @@ def upper_limit(self):
 
 
 class SimplexContinuousDistribution(ContinuousDistribution):
-    ...
+    @property
+    def test_value(self):
+        return tf.ones(self.batch_shape + self.event_shape) / self.event_shape[-1]

pymc4/distributions/multivariate.py

@@ -3,6 +3,7 @@
 Wraps selected tfp.distributions (listed in __all__) as pm.RandomVariables.
 Implements random variables not supported by tfp as distributions.
 """
+import tensorflow as tf
 from tensorflow_probability import distributions as tfd
 from pymc4.distributions.distribution import (
     SimplexContinuousDistribution,
@@ -102,6 +103,10 @@ def _init_distribution(conditions):
         dimension, concentration = conditions["dimension"], conditions["concentration"]
         return tfd.LKJ(dimension=dimension, concentration=concentration)
 
+    @property
+    def test_value(self):
+        return tf.linalg.diag(tf.ones((self.batch_shape + self.event_shape)[:-1]))
+
 
 class Multinomial(DiscreteDistribution):
     r"""
@@ -134,6 +139,8 @@ class Multinomial(DiscreteDistribution):
         Probability of each one of the different outcomes. Elements must
         be non-negative and sum to 1 along the last axis.
     """
+    # For some ridiculous reason, tfp needs multinomial values to be floats...
+    _test_value = 0.0  # type: ignore
 
     def __init__(self, name, total_count, probs, **kwargs):
         super().__init__(name, total_count=total_count, probs=probs, **kwargs)
@@ -268,3 +275,7 @@ def __init__(self, name, df, scale, **kwargs):
     def _init_distribution(conditions):
         df, scale = conditions["df"], conditions["scale"]
         return tfd.WishartTriL(df=df, scale_tril=scale)
+
+    @property
+    def test_value(self):
+        return tf.linalg.diag(tf.ones((self.batch_shape + self.event_shape)[:-1]))