add a draft for traces

pymc4/distributions/tensorflow/continuous.py

@@ -1,6 +1,4 @@
-"""
-PyMC4 continuous random variables for tensorflow.
-"""
+"""PyMC4 continuous random variables for tensorflow."""
 import tensorflow_probability as tfp
 from pymc4.distributions import abstract
 from pymc4.distributions.tensorflow.distribution import BackendDistribution

pymc4/flow/executor.py

@@ -1,5 +1,5 @@
 import types
-from typing import Any, Tuple, Dict, Union, List, Optional
+from typing import Any, Tuple, Dict, Union, List
 import collections
 import itertools
 from pymc4 import _backend
@@ -464,8 +464,7 @@ def modify_distribution(
         return dist
 
     def proceed_distribution(self, dist: abstract.Distribution, state: SamplingState):
-        """TODO
-        """
+        # TODO: docs
         if dist.is_anonymous:
             raise EvaluationError("Attempting to create an anonymous Distribution")
         scoped_name = scopes.variable_name(dist.name)

pymc4/inference/__init__.py

@@ -1,4 +1,5 @@
-from pymc4 import _backend
+from .. import _backend
+from . import utils
 
 if _backend.TENSORFLOW:
     from .tensorflow import *  # pylint: disable=wildcard-import

pymc4/inference/tensorflow/sampling.py

@@ -20,7 +20,7 @@ def sample(
     xla=False,
 ):
     """
-    The main API to perform MCMC sampling using NUTS (for now)
+    Perform MCMC sampling using NUTS (for now).
 
     Parameters
     ----------
@@ -91,14 +91,22 @@ def sample(
 
     """
     logpfn, init = build_logp_function(model, state=state, observed=observed)
-    init = [tf.tile(tf.expand_dims(tens, 0), [num_chains] + [1] * tens.ndim) for tens in init]
-
-    def parallel_logpfn(*state):
-        # NOTE: vmap passes things in a tuple, hack to unwrap
-        return tf.vectorized_map(lambda mini_state: logpfn(*mini_state), state)
+    init_state = list(init.values())
+    init_keys = list(init.keys())
+    parallel_logpfn = vectorize_logp_function(logpfn)
+    init_state = tile_init(init_state, num_chains)
+
+    def trace_fn(_, pkr):
+        return (
+            pkr.inner_results.target_log_prob,
+            pkr.inner_results.leapfrogs_taken,
+            pkr.inner_results.has_divergence,
+            pkr.inner_results.energy,
+            pkr.inner_results.log_accept_ratio
+        )
 
-    @tf.function
-    def run_chains(init):
+    @tf.function(autograph=False)
+    def run_chains(init, step_size):
         nuts_kernel = mcmc.NoUTurnSampler(
             target_log_prob_fn=parallel_logpfn, step_size=step_size, **(nuts_kwargs or dict())
         )
@@ -111,22 +119,26 @@ def run_chains(init):
             **(adaptation_kwargs or dict()),
         )
 
-        results = mcmc.sample_chain(
-            num_samples + burn_in,
+        results, sample_stats = mcmc.sample_chain(
+            num_samples,
             current_state=init,
             kernel=adapt_nuts_kernel,
             num_burnin_steps=burn_in,
+            trace_fn=trace_fn,
             **(sample_chain_kwargs or dict()),
         )
 
-        return results
+        return results, sample_stats
 
     if xla:
-        results, stats = tf.xla.experimental.compile(run_chains, inputs=[init])
+        results, sample_stats = tf.xla.experimental.compile(run_chains, inputs=[init_state, step_size])
     else:
-        results, stats = run_chains(init)
+        results, sample_stats = run_chains(init_state, step_size)
 
-    return results
+    posterior = dict(zip(init_keys, results))
+    # Keep in sync with pymc3 naming convention
+    sampler_stats = dict(zip(['lp', 'tree_size', 'diverging', 'energy', 'mean_tree_accept'], sample_stats))
+    return posterior, sampler_stats
 
 
 def build_logp_function(
@@ -149,11 +161,25 @@ def build_logp_function(
     unobserved_keys, unobserved_values = zip(*state.all_unobserved_values.items())
 
     @tf.function(autograph=False)
-    def logpfn(*values):
-        st = flow.SamplingState.from_values(
-            dict(zip(unobserved_keys, values)), observed_values=observed
-        )
+    def logpfn(*values, **kwargs):
+        if kwargs and values:
+            raise TypeError("Either list state should be passed or a dict one")
+        elif values:
+            kwargs = dict(zip(unobserved_keys, values))
+        st = flow.SamplingState.from_values(kwargs, observed_values=observed)
         _, st = flow.evaluate_model_transformed(model, state=st)
         return st.collect_log_prob()
 
-    return logpfn, list(unobserved_values)
+    return logpfn, dict(state.all_unobserved_values)
+
+
+def vectorize_logp_function(logpfn):
+    # TODO: vectorize with dict
+    def vectorized_logpfn(*state):
+        return tf.vectorized_map(lambda mini_state: logpfn(*mini_state), state)
+
+    return vectorized_logpfn
+
+
+def tile_init(init, num_repeats):
+    return [tf.tile(tf.expand_dims(tens, 0), [num_repeats] + [1] * tens.ndim) for tens in init]

pymc4/inference/utils.py

@@ -1,13 +1,14 @@
 from typing import Optional
+import numpy as np
 
-from pymc4 import Model, flow
+from .. import Model, flow
 
 
 def initialize_state(model: Model, observed: Optional[dict] = None) -> flow.SamplingState:
     """
-    Initilize the model provided state and/or observed variables
+    Initilize the model provided state and/or observed variables.
 
-     Parameters
+    Parameters
     ----------
     model : pymc4.Model
     observed : Optional[dict]
@@ -18,3 +19,26 @@ def initialize_state(model: Model, observed: Optional[dict] = None) -> flow.Samp
     """
     _, state = flow.evaluate_model_transformed(model, observed=observed)
     return state.as_sampling_state()
+
+
+def trace_to_arviz(pm4_trace, pm4_sample_stats):
+    """
+    Tensorflow to Arviz trace convertor.
+
+    Convert a PyMC4 trace as returned by sample() to an ArviZ trace object
+    that can be passed to e.g. arviz.plot_trace().
+
+    Parameters
+    ----------
+    pm4_trace : dict
+
+    Returns
+    -------
+    arviz.data.inference_data.InferenceData
+    """
+    import arviz as az
+
+    posterior = {k: np.swapaxes(v.numpy(), 1, 0) for k, v in pm4_trace.items()}
+    sample_stats = {k: v.numpy().T for k, v in pm4_sample_stats.items()}
+    sample_stats['tree_size'] = np.diff(sample_stats['tree_size'], axis=1)
+    return az.from_dict(posterior=posterior, sample_stats=sample_stats)

tests/test_8schools.py

@@ -9,7 +9,7 @@
 @pm4.model
 def schools_pm4():
     eta = yield pm4.Normal("eta", 0, 1, plate=J)
-    mu = yield pm4.Normal("mu", 1, 1e6)
+    mu = yield pm4.Normal("mu", 1, 10)
     tau = yield pm4.HalfNormal("tau", 1 * 2.0)
 
     theta = mu + tau * eta
@@ -19,8 +19,8 @@ def schools_pm4():
 
 
 def test_sample_no_xla():
-    # TODO: better test, compare to etalon chain from pymc3,
-    #   for now it is only to veryfy it is runnable
+    # TODO: better test, compare to a golden standard chain from pymc3,
+    #   for now it is only to verify it is runnable
     tf_trace = pm4.inference.sampling.sample(
         schools_pm4(), step_size=0.28, num_chains=4, num_samples=100, burn_in=50, xla=False
     )