Repository: arviz-devs/arviz

Issue:
from_tfp with multiple chains?
How should `from_tfp` be used with multiple chains? It looks like it is only compatible with one chain, and it treats the different chains as different variables.

## Example

    import tensorflow as tf
    import tensorflow_probability as tfp
    import tensorflow_probability.python.edward2 as ed
    import numpy as np
    import arviz as az
    dtype = np.float32

    def unnormalized_log_prob(x):
        return -x**2.

    samples_gauss, _ = tfp.mcmc.sample_chain(
        num_results=1000,
        current_state=np.ones(20, dtype=dtype),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            unnormalized_log_prob,
            step_size=1.0,
            num_leapfrog_steps=3),
        num_burnin_steps=500,
        parallel_iterations=4) 

    with tf.Session() as sess:
        [samples_gauss_, ] = sess.run([samples_gauss, ])

    az.from_tfp([samples_gauss_]).posterior 
    # It is expected to be a list with a variable (array) at each position

#### Output
    <xarray.Dataset>
    Dimensions:      (chain: 1, draw: 1000, var_0_dim_0: 20)
    Coordinates:
      * chain        (chain) int64 0
      * draw         (draw) int64 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
      * var_0_dim_0  (var_0_dim_0) int64 0 1 2 3 4 5 6 7 ... 12 13 14 15 16 17 18 19
    Data variables:
        var_0        (chain, draw, var_0_dim_0) float32 0.38988703 ... -0.04516393
    Attributes:
    ...

### Workaround
My array has dimensions (ndraws, nchains), therefore simply transposing the array allows to call:

    from_array = az.convert_to_inference_data(samples_gauss_.T)
    az.plot_trace(from_array);

#### Output
Which as it can be seen, works properly.

![image](https://user-images.githubusercontent.com/23738400/50059636-d9ef5980-0189-11e9-9934-fd144177623e.png)

I don't know if I am missing how to use the coords or dims parameter to somehow sort this out or if it is an implementation detail.



Discussion:
I modified a little the eight schools example from the tests to use multiple chains. The idea is to update the cookbook with this version once I understand how it works. You can find it in my fork's [cookbook](https://github.com/OriolAbril/arviz/blob/cookbook_tfp/doc/notebooks/InferenceDataCookbook.ipynb).
Hi,

So the loop fails in here somehow

```
for i, var_name in enumerate(self.var_names):
            data[var_name] = np.expand_dims(self.posterior[i], axis=0)
```

https://github.com/arviz-devs/arviz/blob/master/arviz/data/io_tfp.py

Is there a different structure for 1 chain vs multiple chains?

`dict_to_dataset` assumes `data.shape == chain, draw, *shape`

https://github.com/arviz-devs/arviz/blob/master/arviz/data/base.py
@ahartikainen any updates on this? Here is an example of running 1 chain vs 3 chains which hopefully answers your question. 

I'm not 100% sure about this but I think it might be hard to automatically tell which dimension represents chains with TFP, I think it depends on how the user has chosen to run multiple chains via extra dimensions in `initial_state`. 

In the very least, the first dimension represents draws in TFP - as far as I can tell:

``` python
import tensorflow as tf
import tensorflow_probability as tfp


def unnormalized_log_prob(x):
    return -x ** 2.0


@tf.function
def sample(initial_state):
    samples = tfp.mcmc.sample_chain(
        num_results=200,
        current_state=[initial_state],
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            unnormalized_log_prob, step_size=1.0, num_leapfrog_steps=3
        ),
        num_burnin_steps=100,
        trace_fn=None,
    )

    return samples


one_chain = sample(tf.ones(1))
three_chains = sample(tf.ones(3))
```

some output:

```
>>> print([x.shape for x in one_chain])
[TensorShape([200, 1])]
>>> print([x.shape for x in three_chains])
[TensorShape([200, 3])]
```
Nothing new yet.

The current master branch has `az.concat` that can be used againts `draw` or `chain` dim to concat multiple InferenceData objects.

For tfp we could add extra kw (chain_idx) which by default could be None or 1 or something suitable. Then users could manually change it. It could even eat dictionary with varnames as keys.
Thanks for the quick reply! I'll check out `az.concat`. Let me know if there is anything I can do to help (not that familiar with `arviz` but I could probably provide some TFP help).  
It would be interesting to see a few different examples with nD (n=2, 3) variable and how increasing chains  (in different positions).

Also is there any heuristics that could be used to solve the location and num of chains.

If you want, you could try to fix this with a PR. I think together with arviz team we can make it work.
From @junpenglao in regards to batch shape


> The problem I think is that it is often unclear what is the batch shape, which likely need user input
> for example you can have RV_1_trace.shape = [nsamples, batch_shape1, batch_shape2, event_shape....]
@junpenglao @jeffpollock9 What about adding 2 arguments: `chain_dim` and `draw_dim`, then, if `chain_dim` is `None`, a new dimension is added at the beggining of the array, otherwise, `chain_dim` is moved to be the first one and `draw_dim` the second. 

This could even be reduced to only one new argument dropping `draw_dim` if the first dimension of tensorflow results always corresponds to draws.
That should work if nobody is doing funny thing like chain_shape=(2,5).
Should we support multiple chain dims? We could stack them before conversion to dataset
I am intend to say not supporting, I haven't really seen any valid use case.

Resolving patch:
diff --git a/arviz/data/io_tfp.py b/arviz/data/io_tfp.py
--- a/arviz/data/io_tfp.py
+++ b/arviz/data/io_tfp.py
@@ -19,6 +19,7 @@ def __init__(
         feed_dict=None,
         posterior_predictive_samples=100,
         posterior_predictive_size=1,
+        chain_dim=None,
         observed=None,
         coords=None,
         dims=None
@@ -38,6 +39,7 @@ def __init__(
         self.posterior_predictive_samples = posterior_predictive_samples
         self.posterior_predictive_size = posterior_predictive_size
         self.observed = observed
+        self.chain_dim = chain_dim
         self.coords = coords
         self.dims = dims
 
@@ -55,11 +57,20 @@ def __init__(
             tf.disable_v2_behavior()
             self.tf = tf  # pylint: disable=invalid-name
 
+    def handle_chain_location(self, ary):
+        """Move the axis corresponding to the chain to first position.
+
+        If there is only one chain which has no axis, add it.
+        """
+        if self.chain_dim is None:
+            return utils.expand_dims(ary)
+        return ary.swapaxes(0, self.chain_dim)
+
     def posterior_to_xarray(self):
         """Convert the posterior to an xarray dataset."""
         data = {}
         for i, var_name in enumerate(self.var_names):
-            data[var_name] = utils.expand_dims(self.posterior[i])
+            data[var_name] = self.handle_chain_location(self.posterior[i])
         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)
 
     def observed_data_to_xarray(self):
@@ -121,7 +132,9 @@ def posterior_predictive_to_xarray(self):
 
         data = {}
         with self.tf.Session() as sess:
-            data["obs"] = utils.expand_dims(sess.run(posterior_preds, feed_dict=self.feed_dict))
+            data["obs"] = self.handle_chain_location(
+                sess.run(posterior_preds, feed_dict=self.feed_dict)
+            )
         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)
 
     def sample_stats_to_xarray(self):
@@ -148,7 +161,7 @@ def sample_stats_to_xarray(self):
         dims = {"log_likelihood": coord_name}
 
         with self.tf.Session() as sess:
-            data["log_likelihood"] = utils.expand_dims(
+            data["log_likelihood"] = self.handle_chain_location(
                 sess.run(log_likelihood, feed_dict=self.feed_dict)
             )
         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=dims)
@@ -178,6 +191,7 @@ def from_tfp(
     feed_dict=None,
     posterior_predictive_samples=100,
     posterior_predictive_size=1,
+    chain_dim=None,
     observed=None,
     coords=None,
     dims=None
@@ -190,6 +204,7 @@ def from_tfp(
         feed_dict=feed_dict,
         posterior_predictive_samples=posterior_predictive_samples,
         posterior_predictive_size=posterior_predictive_size,
+        chain_dim=chain_dim,
         observed=observed,
         coords=coords,
         dims=dims,


Test patch: