Compound/Gibbs step and discrete suppot

nbconfig.py

@@ -15,7 +15,6 @@
     "notebooks/baseball.ipynb",
     "notebooks/basic-usage.ipynb",
     "notebooks/rugby_analytics.ipynb",
-# will reinstate in a later PR
-#     "notebooks/radon_hierarchical.ipynb",
+    # will reinstate in a later PR
+    #     "notebooks/radon_hierarchical.ipynb",
 ]
-

notebooks/utils.py

@@ -1,15 +1,16 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+
 def plot_samples(x, batched_samples, labels, names, ylim=None):
     if not isinstance(batched_samples, np.ndarray):
         batched_samples = np.asarray(batched_samples)
     n_samples = batched_samples.shape[0]
     if ylim is not None:
         ymin, ymax = ylim
     else:
-        ymin, ymax = batched_samples.min()-0.2, batched_samples.max()+0.2
-    fig, ax = plt.subplots(n_samples, 1, figsize=(14, n_samples*3))
+        ymin, ymax = batched_samples.min() - 0.2, batched_samples.max() + 0.2
+    fig, ax = plt.subplots(n_samples, 1, figsize=(14, n_samples * 3))
     if isinstance(labels, (list, tuple)):
         labels = [np.asarray(label) for label in labels]
     else:
@@ -29,6 +30,7 @@ def plot_samples(x, batched_samples, labels, names, ylim=None):
         axi.set_title(lab)
     plt.show()
 
+
 def plot_cov_matrix(k, X, labels, names, vlim=None, cmap="inferno", interpolation="none"):
     cov = k(X, X)
     cov = np.asarray(cov)
@@ -42,8 +44,9 @@ def plot_cov_matrix(k, X, labels, names, vlim=None, cmap="inferno", interpolatio
     else:
         labels = np.asarray(labels)
         n_samples = 1
-    fig, ax = plt.subplots(1, n_samples, figsize=(5*n_samples, 4))
-    if not isinstance(ax, np.ndarray): ax = np.asarray([ax])
+    fig, ax = plt.subplots(1, n_samples, figsize=(5 * n_samples, 4))
+    if not isinstance(ax, np.ndarray):
+        ax = np.asarray([ax])
     for i in range(ax.shape[0]):
         axi = ax[i]
         if isinstance(labels, (list, tuple)):

pymc4/__init__.py

@@ -1,6 +1,6 @@
 """PyMC4."""
-
 from . import utils
+from .coroutine_model import Model, model
 from .scopes import name_scope, variable_name
 from . import coroutine_model
 from . import distributions
@@ -12,12 +12,14 @@
     evaluate_meta_model,
     evaluate_meta_posterior_predictive_model,
 )
-from .coroutine_model import Model, model
 from . import inference
 from .distributions import *
 from .forward_sampling import sample_prior_predictive, sample_posterior_predictive
 from .inference.sampling import sample
+from .mcmc.samplers import *
 from . import gp
+from . import mcmc
 from .variational import *
 
+
 __version__ = "4.0a2"

pymc4/distributions/__init__.py

@@ -5,3 +5,5 @@
 from .distribution import Potential, Deterministic
 from .mixture import Mixture
 from . import transforms
+from .mixture import *
+from .state_functions import *

pymc4/distributions/batchstack.py

@@ -30,7 +30,7 @@ def _fn(self, **kwargs):
         shape = prefer_static.concat(
             [
                 prefer_static.ones(
-                    prefer_static.rank_from_shape(self.batch_stack), dtype=self.batch_stack.dtype
+                    prefer_static.rank_from_shape(self.batch_stack), dtype=self.batch_stack.dtype,
                 ),
                 self.distribution.batch_shape_tensor(),
                 self.distribution.event_shape_tensor(),
@@ -60,7 +60,7 @@ class BatchStacker(distribution_lib.Distribution):
 
     The probability function is,
 
-    .. math:: 
+    .. math::
         p(x) = prod{ p(x[i]) : i = 0, ..., (n - 1) }
 
     Examples
@@ -77,7 +77,7 @@ class BatchStacker(distribution_lib.Distribution):
     >>> lp = s.log_prob(x)
     >>> lp.shape.as_list()
     [5]
-    
+
     Example 2: `[5, 4]`-draws of a bivariate Normal.
 
     >>> s = BatchStacker(
@@ -186,7 +186,7 @@ def _log_prob(self, x, **kwargs):
         x = tf.reshape(
             x,
             shape=tf.pad(
-                tf.shape(x), paddings=[[prefer_static.maximum(0, -d), 0]], constant_values=1
+                tf.shape(x), paddings=[[prefer_static.maximum(0, -d), 0]], constant_values=1,
             ),
         )
         # (2) Compute x's log_prob.

pymc4/distributions/continuous.py

@@ -120,7 +120,7 @@ class GeneralizedNormal(ContinuousDistribution):
 
     .. math::
 
-       f(x \mid \mu, \alpha, \beta) = 
+       f(x \mid \mu, \alpha, \beta) =
            \frac{\beta}{2 \Gamma(1/\beta)}
            \exp(-(|x - \mu| /\alpha)^\beta)
 
@@ -354,7 +354,10 @@ def __init__(self, name, concentration0, concentration1, **kwargs):
 
     @staticmethod
     def _init_distribution(conditions, **kwargs):
-        concentration0, concentration1 = conditions["concentration0"], conditions["concentration1"]
+        concentration0, concentration1 = (
+            conditions["concentration0"],
+            conditions["concentration1"],
+        )
         return tfd.Beta(concentration0=concentration0, concentration1=concentration1, **kwargs)
 
 
@@ -779,7 +782,10 @@ def __init__(self, name, concentration0, concentration1, **kwargs):
 
     @staticmethod
     def _init_distribution(conditions, **kwargs):
-        concentration0, concentration1 = conditions["concentration0"], conditions["concentration1"]
+        concentration0, concentration1 = (
+            conditions["concentration0"],
+            conditions["concentration1"],
+        )
         return tfd.Kumaraswamy(
             concentration0=concentration0, concentration1=concentration1, **kwargs
         )
@@ -987,7 +993,7 @@ class Moyal(ContinuousDistribution):
 
     .. math::
 
-        f(x \mid \mu, \sigma) = 
+        f(x \mid \mu, \sigma) =
            \frac{1}{\sqrt{2\pi}\sigma}
            \exp\left(-\frac{1}{2}\left[\frac{x-\mu}{\sigma}+\exp\left(-\frac{x-\mu}{\sigma}\right)\right]\right)
 
@@ -1493,6 +1499,13 @@ class Weibull(PositiveContinuousDistribution):
         Shape parameter (concentration > 0).
     scale : float|tensor
         Scale parameter (scale > 0).
+
+    Developer Notes
+    ---------------
+    The Weibull distribution is implemented as a standard uniform distribution transformed by the
+    Inverse of the WeibullCDF bijector. The shape to broadcast the low and high parameters for the
+    Uniform distribution are obtained using
+    tensorflow_probability.python.internal.distribution_util.prefer_static_broadcast_shape()
     """
 
     def __init__(self, name, concentration, scale, **kwargs):

pymc4/distributions/discrete.py

@@ -1,12 +1,18 @@
 """PyMC4 discrete random variables."""
 import tensorflow as tf
 from tensorflow_probability import distributions as tfd
+from tensorflow_probability.python.internal import prefer_static
 from pymc4.distributions.distribution import (
     PositiveDiscreteDistribution,
     BoundedDiscreteDistribution,
 )
 from pymc4.distributions import transforms
 
+from pymc4.distributions.state_functions import (
+    categorical_uniform_fn,
+    bernoulli_fn,
+)
+
 __all__ = [
     "Bernoulli",
     "Binomial",
@@ -58,9 +64,11 @@ class Bernoulli(BoundedDiscreteDistribution):
     probs : float
         Probability of success (0 < probs < 1).
     """
+    _grad_support = False
 
     def __init__(self, name, probs, **kwargs):
         super().__init__(name, probs=probs, **kwargs)
+        self._default_new_state_part = bernoulli_fn()
 
     @staticmethod
     def _init_distribution(conditions, **kwargs):
@@ -245,6 +253,7 @@ class DiscreteUniform(BoundedDiscreteDistribution):
     high : int
         Upper limit (high > low).
     """
+    _grad_support = False
 
     def __init__(self, name, low, high, **kwargs):
         super().__init__(name, low=low, high=high, **kwargs)
@@ -296,9 +305,12 @@ class Categorical(BoundedDiscreteDistribution):
     probs : array of floats
         probs > 0 and the elements of probs must sum to 1.
     """
+    _grad_support = False
 
     def __init__(self, name, probs, **kwargs):
         super().__init__(name, probs=probs, **kwargs)
+        classes = prefer_static.shape(probs)[-1]
+        self._default_new_state_part = categorical_uniform_fn(classes=classes)
 
     @staticmethod
     def _init_distribution(conditions, **kwargs):
@@ -310,7 +322,7 @@ def lower_limit(self):
         return 0.0
 
     def upper_limit(self):
-        return self.conditions["probs"].shape[-1]
+        return float(tf.shape(self.conditions["probs"])[-1])
 
 
 class Geometric(BoundedDiscreteDistribution):
@@ -486,8 +498,6 @@ class Poisson(PositiveDiscreteDistribution):
     """
 
     # 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)
 
@@ -770,4 +780,4 @@ def lower_limit(self):
         return 0.0
 
     def upper_limit(self):
-        return self.conditions["cutpoints"].shape[-1]
+        return prefer_static.shape(self.conditions["cutpoints"])[-1]

pymc4/distributions/distribution.py

@@ -7,7 +7,7 @@
 from tensorflow_probability import distributions as tfd
 from pymc4.coroutine_model import Model, unpack
 from pymc4.distributions.batchstack import BatchStacker
-from . import transforms
+from pymc4.distributions import transforms
 
 NameType = Union[str, int]
 
@@ -29,6 +29,7 @@
 class Distribution(Model):
     """Statistical distribution."""
 
+    _grad_support: bool = True
     _test_value = 0.0
     _base_parameters = ["dtype", "validate_args", "allow_nan_stats"]
 
@@ -48,9 +49,10 @@ def __init__(
         **kwargs,
     ):
         self.conditions, self.base_parameters = self.unpack_conditions(
-            dtype=dtype, validate_args=validate_args, allow_nan_stats=allow_nan_stats, **kwargs
+            dtype=dtype, validate_args=validate_args, allow_nan_stats=allow_nan_stats, **kwargs,
         )
         self._distribution = self._init_distribution(self.conditions, **self.base_parameters)
+        self._default_new_state_part = None
         super().__init__(
             self.unpack_distribution, name=name, keep_return=True, keep_auxiliary=False
         )
@@ -104,7 +106,7 @@ def unpack_conditions(cls, **kwargs) -> Tuple[dict, dict]:
     @property
     def test_value(self):
         return tf.cast(
-            tf.broadcast_to(self._test_value, self.batch_shape + self.event_shape), self.dtype
+            tf.broadcast_to(self._test_value, self.batch_shape + self.event_shape), self.dtype,
         )
 
     def sample(self, sample_shape=(), seed=None):
@@ -139,14 +141,14 @@ def sample_numpy(self, sample_shape=(), seed=None):
     def get_test_sample(self, sample_shape=(), seed=None):
         """
         Get the test value using a function signature similar to meth:`~.sample`.
-        
+
         Parameters
         ----------
         sample_shape : tuple
             sample shape
         seed : int | None
             ignored. Is only present to match the signature of meth:`~.sample`
-        
+
         Returns
         -------
         The distribution's ``test_value`` broadcasted to
@@ -295,6 +297,9 @@ def upper_limit(self):
 
 
 class BoundedDiscreteDistribution(DiscreteDistribution, BoundedDistribution):
+    def _init_transform(self, transform):
+        return transform
+
     @property
     def _test_value(self):
         return tf.cast(tf.round(0.5 * (self.upper_limit() + self.lower_limit())), self.dtype)
@@ -339,12 +344,6 @@ def upper_limit(self):
 class PositiveDiscreteDistribution(BoundedDiscreteDistribution):
     _test_value = 1
 
-    def _init_transform(self, transform):
-        if transform is None:
-            return transforms.Log()
-        else:
-            return transform
-
     def lower_limit(self):
         return 0
 

pymc4/distributions/half_student_t.py

@@ -40,7 +40,7 @@ class HalfStudentT(distribution.Distribution):
     """
 
     def __init__(
-        self, df, loc, scale, validate_args=False, allow_nan_stats=True, name="HalfStudentT"
+        self, df, loc, scale, validate_args=False, allow_nan_stats=True, name="HalfStudentT",
     ):
         r"""
         Construct a half-Student's t distribution with ``df``, ``loc`` and ``scale``.
@@ -93,7 +93,7 @@ def __init__(
     @staticmethod
     def _param_shapes(sample_shape):
         return dict(
-            zip(("df", "loc", "scale"), ([tf.convert_to_tensor(sample_shape, dtype=tf.int32)] * 3))
+            zip(("df", "loc", "scale"), ([tf.convert_to_tensor(sample_shape, dtype=tf.int32)] * 3),)
         )
 
     @classmethod
@@ -252,7 +252,7 @@ def _variance(self):
         )
         if self.allow_nan_stats:
             return tf.where(
-                df > 1.0, result_where_defined, dtype_util.as_numpy_dtype(self.dtype)(np.nan)
+                df > 1.0, result_where_defined, dtype_util.as_numpy_dtype(self.dtype)(np.nan),
             )
         else:
             return distribution_util.with_dependencies(

pymc4/distributions/mixture.py

@@ -15,16 +15,12 @@
 class Mixture(Distribution):
     r"""
     Mixture random variable.
-
     Often used to model subpopulation heterogeneity
-
     .. math:: f(x \mid w, \theta) = \sum_{i = 1}^n w_i f_i(x \mid \theta_i)
-
     ========  ============================================
     Support   :math:`\cap_{i = 1}^n \textrm{support}(f_i)`
     Mean      :math:`\sum_{i = 1}^n w_i \mu_i`
     ========  ============================================
-
     Parameters
     ----------
     p : tf.Tensor
@@ -98,7 +94,7 @@ def _init_distribution(conditions, **kwargs):
                 )
             distr = [el._distribution for el in d]
             return tfd.Mixture(
-                tfd.Categorical(probs=p, **kwargs), distr, **kwargs, use_static_graph=True
+                tfd.Categorical(probs=p, **kwargs), distr, **kwargs, use_static_graph=True,
             )
         # else if 'd' is a pymc distribution with batch_size > 1
         elif isinstance(d, Distribution):

pymc4/distributions/multivariate.py

@@ -233,7 +233,10 @@ def __init__(self, name, mean_direction, concentration, **kwargs):
 
     @staticmethod
     def _init_distribution(conditions, **kwargs):
-        mean_direction, concentration = conditions["mean_direction"], conditions["concentration"]
+        mean_direction, concentration = (
+            conditions["mean_direction"],
+            conditions["concentration"],
+        )
         return tfd.VonMisesFisher(
             mean_direction=mean_direction, concentration=concentration, **kwargs
         )