We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `arviz/data/io_tfp.py`
Content:
```
1 """Tfp-specific conversion code."""
2 import numpy as np
3 import xarray as xr
4 
5 from .inference_data import InferenceData
6 from .base import dict_to_dataset, generate_dims_coords, make_attrs
7 from .. import utils
8 
9 # pylint: disable=too-many-instance-attributes
10 class TfpConverter:
11     """Encapsulate tfp specific logic."""
12 
13     def __init__(
14         self,
15         *,
16         posterior,
17         var_names=None,
18         model_fn=None,
19         feed_dict=None,
20         posterior_predictive_samples=100,
21         posterior_predictive_size=1,
22         observed=None,
23         coords=None,
24         dims=None
25     ):
26 
27         self.posterior = posterior
28 
29         if var_names is None:
30             self.var_names = []
31             for i in range(0, len(posterior)):
32                 self.var_names.append("var_{0}".format(i))
33         else:
34             self.var_names = var_names
35 
36         self.model_fn = model_fn
37         self.feed_dict = feed_dict
38         self.posterior_predictive_samples = posterior_predictive_samples
39         self.posterior_predictive_size = posterior_predictive_size
40         self.observed = observed
41         self.coords = coords
42         self.dims = dims
43 
44         import tensorflow_probability as tfp
45         import tensorflow as tf
46         import tensorflow_probability.python.edward2 as ed
47 
48         self.tfp = tfp
49         self.tf = tf  # pylint: disable=invalid-name
50         self.ed = ed  # pylint: disable=invalid-name
51 
52         if int(self.tf.__version__[0]) > 1:
53             import tensorflow.compat.v1 as tf  # pylint: disable=import-error
54 
55             tf.disable_v2_behavior()
56             self.tf = tf  # pylint: disable=invalid-name
57 
58     def posterior_to_xarray(self):
59         """Convert the posterior to an xarray dataset."""
60         data = {}
61         for i, var_name in enumerate(self.var_names):
62             data[var_name] = utils.expand_dims(self.posterior[i])
63         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)
64 
65     def observed_data_to_xarray(self):
66         """Convert observed data to xarray."""
67         if self.observed is None:
68             return None
69 
70         observed_data = {}
71         if isinstance(self.observed, self.tf.Tensor):
72             with self.tf.Session() as sess:
73                 vals = sess.run(self.observed, feed_dict=self.feed_dict)
74         else:
75             vals = self.observed
76 
77         if self.dims is None:
78             dims = {}
79         else:
80             dims = self.dims
81 
82         name = "obs"
83         val_dims = dims.get(name)
84         vals = utils.one_de(vals)
85         val_dims, coords = generate_dims_coords(vals.shape, name, dims=val_dims, coords=self.coords)
86         # coords = {key: xr.IndexVariable((key,), data=coords[key]) for key in val_dims}
87 
88         observed_data[name] = xr.DataArray(vals, dims=val_dims, coords=coords)
89         return xr.Dataset(data_vars=observed_data, attrs=make_attrs(library=self.tfp))
90 
91     def _value_setter(self, variables):
92         def interceptor(rv_constructor, *rv_args, **rv_kwargs):
93             """Replace prior on effects with empirical posterior mean from MCMC."""
94             name = rv_kwargs.pop("name")
95             if name in variables:
96                 rv_kwargs["value"] = variables[name]
97             return rv_constructor(*rv_args, **rv_kwargs)
98 
99         return interceptor
100 
101     def posterior_predictive_to_xarray(self):
102         """Convert posterior_predictive samples to xarray."""
103         if self.model_fn is None:
104             return None
105 
106         posterior_preds = []
107         sample_size = self.posterior[0].shape[0]
108 
109         for i in np.arange(0, sample_size, int(sample_size / self.posterior_predictive_samples)):
110             variables = {}
111             for var_i, var_name in enumerate(self.var_names):
112                 variables[var_name] = self.posterior[var_i][i]
113 
114             with self.ed.interception(self._value_setter(variables)):
115                 if self.posterior_predictive_size > 1:
116                     posterior_preds.append(
117                         [self.model_fn() for _ in range(self.posterior_predictive_size)]
118                     )
119                 else:
120                     posterior_preds.append(self.model_fn())
121 
122         data = {}
123         with self.tf.Session() as sess:
124             data["obs"] = utils.expand_dims(sess.run(posterior_preds, feed_dict=self.feed_dict))
125         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)
126 
127     def sample_stats_to_xarray(self):
128         """Extract sample_stats from tfp trace."""
129         if self.model_fn is None or self.observed is None:
130             return None
131 
132         log_likelihood = []
133         sample_size = self.posterior[0].shape[0]
134 
135         for i in range(sample_size):
136             variables = {}
137             for var_i, var_name in enumerate(self.var_names):
138                 variables[var_name] = self.posterior[var_i][i]
139 
140             with self.ed.interception(self._value_setter(variables)):
141                 log_likelihood.append((self.model_fn().distribution.log_prob(self.observed)))
142 
143         data = {}
144         if self.dims is not None:
145             coord_name = self.dims.get("obs")
146         else:
147             coord_name = None
148         dims = {"log_likelihood": coord_name}
149 
150         with self.tf.Session() as sess:
151             data["log_likelihood"] = utils.expand_dims(
152                 sess.run(log_likelihood, feed_dict=self.feed_dict)
153             )
154         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=dims)
155 
156     def to_inference_data(self):
157         """Convert all available data to an InferenceData object.
158 
159         Note that if groups can not be created (i.e., there is no `trace`, so
160         the `posterior` and `sample_stats` can not be extracted), then the InferenceData
161         will not have those groups.
162         """
163         return InferenceData(
164             **{
165                 "posterior": self.posterior_to_xarray(),
166                 "sample_stats": self.sample_stats_to_xarray(),
167                 "posterior_predictive": self.posterior_predictive_to_xarray(),
168                 "observed_data": self.observed_data_to_xarray(),
169             }
170         )
171 
172 
173 def from_tfp(
174     posterior=None,
175     *,
176     var_names=None,
177     model_fn=None,
178     feed_dict=None,
179     posterior_predictive_samples=100,
180     posterior_predictive_size=1,
181     observed=None,
182     coords=None,
183     dims=None
184 ):
185     """Convert tfp data into an InferenceData object."""
186     return TfpConverter(
187         posterior=posterior,
188         var_names=var_names,
189         model_fn=model_fn,
190         feed_dict=feed_dict,
191         posterior_predictive_samples=posterior_predictive_samples,
192         posterior_predictive_size=posterior_predictive_size,
193         observed=observed,
194         coords=coords,
195         dims=dims,
196     ).to_inference_data()
197
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         feed_dict=None,
         posterior_predictive_samples=100,
         posterior_predictive_size=1,
+        chain_dim=None,
         observed=None,
         coords=None,
         dims=None
@@ -38,6 +39,7 @@
         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 @@
             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 @@
 
         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 @@
         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 @@
     feed_dict=None,
     posterior_predictive_samples=100,
     posterior_predictive_size=1,
+    chain_dim=None,
     observed=None,
     coords=None,
     dims=None
@@ -190,6 +204,7 @@
         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,

{"golden_diff": "diff --git a/arviz/data/io_tfp.py b/arviz/data/io_tfp.py\n--- a/arviz/data/io_tfp.py\n+++ b/arviz/data/io_tfp.py\n@@ -19,6 +19,7 @@\n         feed_dict=None,\r\n         posterior_predictive_samples=100,\r\n         posterior_predictive_size=1,\r\n+        chain_dim=None,\r\n         observed=None,\r\n         coords=None,\r\n         dims=None\r\n@@ -38,6 +39,7 @@\n         self.posterior_predictive_samples = posterior_predictive_samples\r\n         self.posterior_predictive_size = posterior_predictive_size\r\n         self.observed = observed\r\n+        self.chain_dim = chain_dim\r\n         self.coords = coords\r\n         self.dims = dims\r\n \r\n@@ -55,11 +57,20 @@\n             tf.disable_v2_behavior()\r\n             self.tf = tf  # pylint: disable=invalid-name\r\n \r\n+    def handle_chain_location(self, ary):\r\n+        \"\"\"Move the axis corresponding to the chain to first position.\r\n+\r\n+        If there is only one chain which has no axis, add it.\r\n+        \"\"\"\r\n+        if self.chain_dim is None:\r\n+            return utils.expand_dims(ary)\r\n+        return ary.swapaxes(0, self.chain_dim)\r\n+\r\n     def posterior_to_xarray(self):\r\n         \"\"\"Convert the posterior to an xarray dataset.\"\"\"\r\n         data = {}\r\n         for i, var_name in enumerate(self.var_names):\r\n-            data[var_name] = utils.expand_dims(self.posterior[i])\r\n+            data[var_name] = self.handle_chain_location(self.posterior[i])\r\n         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n \r\n     def observed_data_to_xarray(self):\r\n@@ -121,7 +132,9 @@\n \r\n         data = {}\r\n         with self.tf.Session() as sess:\r\n-            data[\"obs\"] = utils.expand_dims(sess.run(posterior_preds, feed_dict=self.feed_dict))\r\n+            data[\"obs\"] = self.handle_chain_location(\r\n+                sess.run(posterior_preds, feed_dict=self.feed_dict)\r\n+            )\r\n         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n \r\n     def sample_stats_to_xarray(self):\r\n@@ -148,7 +161,7 @@\n         dims = {\"log_likelihood\": coord_name}\r\n \r\n         with self.tf.Session() as sess:\r\n-            data[\"log_likelihood\"] = utils.expand_dims(\r\n+            data[\"log_likelihood\"] = self.handle_chain_location(\r\n                 sess.run(log_likelihood, feed_dict=self.feed_dict)\r\n             )\r\n         return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=dims)\r\n@@ -178,6 +191,7 @@\n     feed_dict=None,\r\n     posterior_predictive_samples=100,\r\n     posterior_predictive_size=1,\r\n+    chain_dim=None,\r\n     observed=None,\r\n     coords=None,\r\n     dims=None\r\n@@ -190,6 +204,7 @@\n         feed_dict=feed_dict,\r\n         posterior_predictive_samples=posterior_predictive_samples,\r\n         posterior_predictive_size=posterior_predictive_size,\r\n+        chain_dim=chain_dim,\r\n         observed=observed,\r\n         coords=coords,\r\n         dims=dims,\n", "issue": "from_tfp with multiple chains?\nHow 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.\r\n\r\n## Example\r\n\r\n    import tensorflow as tf\r\n    import tensorflow_probability as tfp\r\n    import tensorflow_probability.python.edward2 as ed\r\n    import numpy as np\r\n    import arviz as az\r\n    dtype = np.float32\r\n\r\n    def unnormalized_log_prob(x):\r\n        return -x**2.\r\n\r\n    samples_gauss, _ = tfp.mcmc.sample_chain(\r\n        num_results=1000,\r\n        current_state=np.ones(20, dtype=dtype),\r\n        kernel=tfp.mcmc.HamiltonianMonteCarlo(\r\n            unnormalized_log_prob,\r\n            step_size=1.0,\r\n            num_leapfrog_steps=3),\r\n        num_burnin_steps=500,\r\n        parallel_iterations=4) \r\n\r\n    with tf.Session() as sess:\r\n        [samples_gauss_, ] = sess.run([samples_gauss, ])\r\n\r\n    az.from_tfp([samples_gauss_]).posterior \r\n    # It is expected to be a list with a variable (array) at each position\r\n\r\n#### Output\r\n    <xarray.Dataset>\r\n    Dimensions:      (chain: 1, draw: 1000, var_0_dim_0: 20)\r\n    Coordinates:\r\n      * chain        (chain) int64 0\r\n      * draw         (draw) int64 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999\r\n      * 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\r\n    Data variables:\r\n        var_0        (chain, draw, var_0_dim_0) float32 0.38988703 ... -0.04516393\r\n    Attributes:\r\n    ...\r\n\r\n### Workaround\r\nMy array has dimensions (ndraws, nchains), therefore simply transposing the array allows to call:\r\n\r\n    from_array = az.convert_to_inference_data(samples_gauss_.T)\r\n    az.plot_trace(from_array);\r\n\r\n#### Output\r\nWhich as it can be seen, works properly.\r\n\r\n![image](https://user-images.githubusercontent.com/23738400/50059636-d9ef5980-0189-11e9-9934-fd144177623e.png)\r\n\r\nI 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.\r\n\n", "before_files": [{"content": "\"\"\"Tfp-specific conversion code.\"\"\"\r\nimport numpy as np\r\nimport xarray as xr\r\n\r\nfrom .inference_data import InferenceData\r\nfrom .base import dict_to_dataset, generate_dims_coords, make_attrs\r\nfrom .. import utils\r\n\r\n# pylint: disable=too-many-instance-attributes\r\nclass TfpConverter:\r\n    \"\"\"Encapsulate tfp specific logic.\"\"\"\r\n\r\n    def __init__(\r\n        self,\r\n        *,\r\n        posterior,\r\n        var_names=None,\r\n        model_fn=None,\r\n        feed_dict=None,\r\n        posterior_predictive_samples=100,\r\n        posterior_predictive_size=1,\r\n        observed=None,\r\n        coords=None,\r\n        dims=None\r\n    ):\r\n\r\n        self.posterior = posterior\r\n\r\n        if var_names is None:\r\n            self.var_names = []\r\n            for i in range(0, len(posterior)):\r\n                self.var_names.append(\"var_{0}\".format(i))\r\n        else:\r\n            self.var_names = var_names\r\n\r\n        self.model_fn = model_fn\r\n        self.feed_dict = feed_dict\r\n        self.posterior_predictive_samples = posterior_predictive_samples\r\n        self.posterior_predictive_size = posterior_predictive_size\r\n        self.observed = observed\r\n        self.coords = coords\r\n        self.dims = dims\r\n\r\n        import tensorflow_probability as tfp\r\n        import tensorflow as tf\r\n        import tensorflow_probability.python.edward2 as ed\r\n\r\n        self.tfp = tfp\r\n        self.tf = tf  # pylint: disable=invalid-name\r\n        self.ed = ed  # pylint: disable=invalid-name\r\n\r\n        if int(self.tf.__version__[0]) > 1:\r\n            import tensorflow.compat.v1 as tf  # pylint: disable=import-error\r\n\r\n            tf.disable_v2_behavior()\r\n            self.tf = tf  # pylint: disable=invalid-name\r\n\r\n    def posterior_to_xarray(self):\r\n        \"\"\"Convert the posterior to an xarray dataset.\"\"\"\r\n        data = {}\r\n        for i, var_name in enumerate(self.var_names):\r\n            data[var_name] = utils.expand_dims(self.posterior[i])\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n\r\n    def observed_data_to_xarray(self):\r\n        \"\"\"Convert observed data to xarray.\"\"\"\r\n        if self.observed is None:\r\n            return None\r\n\r\n        observed_data = {}\r\n        if isinstance(self.observed, self.tf.Tensor):\r\n            with self.tf.Session() as sess:\r\n                vals = sess.run(self.observed, feed_dict=self.feed_dict)\r\n        else:\r\n            vals = self.observed\r\n\r\n        if self.dims is None:\r\n            dims = {}\r\n        else:\r\n            dims = self.dims\r\n\r\n        name = \"obs\"\r\n        val_dims = dims.get(name)\r\n        vals = utils.one_de(vals)\r\n        val_dims, coords = generate_dims_coords(vals.shape, name, dims=val_dims, coords=self.coords)\r\n        # coords = {key: xr.IndexVariable((key,), data=coords[key]) for key in val_dims}\r\n\r\n        observed_data[name] = xr.DataArray(vals, dims=val_dims, coords=coords)\r\n        return xr.Dataset(data_vars=observed_data, attrs=make_attrs(library=self.tfp))\r\n\r\n    def _value_setter(self, variables):\r\n        def interceptor(rv_constructor, *rv_args, **rv_kwargs):\r\n            \"\"\"Replace prior on effects with empirical posterior mean from MCMC.\"\"\"\r\n            name = rv_kwargs.pop(\"name\")\r\n            if name in variables:\r\n                rv_kwargs[\"value\"] = variables[name]\r\n            return rv_constructor(*rv_args, **rv_kwargs)\r\n\r\n        return interceptor\r\n\r\n    def posterior_predictive_to_xarray(self):\r\n        \"\"\"Convert posterior_predictive samples to xarray.\"\"\"\r\n        if self.model_fn is None:\r\n            return None\r\n\r\n        posterior_preds = []\r\n        sample_size = self.posterior[0].shape[0]\r\n\r\n        for i in np.arange(0, sample_size, int(sample_size / self.posterior_predictive_samples)):\r\n            variables = {}\r\n            for var_i, var_name in enumerate(self.var_names):\r\n                variables[var_name] = self.posterior[var_i][i]\r\n\r\n            with self.ed.interception(self._value_setter(variables)):\r\n                if self.posterior_predictive_size > 1:\r\n                    posterior_preds.append(\r\n                        [self.model_fn() for _ in range(self.posterior_predictive_size)]\r\n                    )\r\n                else:\r\n                    posterior_preds.append(self.model_fn())\r\n\r\n        data = {}\r\n        with self.tf.Session() as sess:\r\n            data[\"obs\"] = utils.expand_dims(sess.run(posterior_preds, feed_dict=self.feed_dict))\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n\r\n    def sample_stats_to_xarray(self):\r\n        \"\"\"Extract sample_stats from tfp trace.\"\"\"\r\n        if self.model_fn is None or self.observed is None:\r\n            return None\r\n\r\n        log_likelihood = []\r\n        sample_size = self.posterior[0].shape[0]\r\n\r\n        for i in range(sample_size):\r\n            variables = {}\r\n            for var_i, var_name in enumerate(self.var_names):\r\n                variables[var_name] = self.posterior[var_i][i]\r\n\r\n            with self.ed.interception(self._value_setter(variables)):\r\n                log_likelihood.append((self.model_fn().distribution.log_prob(self.observed)))\r\n\r\n        data = {}\r\n        if self.dims is not None:\r\n            coord_name = self.dims.get(\"obs\")\r\n        else:\r\n            coord_name = None\r\n        dims = {\"log_likelihood\": coord_name}\r\n\r\n        with self.tf.Session() as sess:\r\n            data[\"log_likelihood\"] = utils.expand_dims(\r\n                sess.run(log_likelihood, feed_dict=self.feed_dict)\r\n            )\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=dims)\r\n\r\n    def to_inference_data(self):\r\n        \"\"\"Convert all available data to an InferenceData object.\r\n\r\n        Note that if groups can not be created (i.e., there is no `trace`, so\r\n        the `posterior` and `sample_stats` can not be extracted), then the InferenceData\r\n        will not have those groups.\r\n        \"\"\"\r\n        return InferenceData(\r\n            **{\r\n                \"posterior\": self.posterior_to_xarray(),\r\n                \"sample_stats\": self.sample_stats_to_xarray(),\r\n                \"posterior_predictive\": self.posterior_predictive_to_xarray(),\r\n                \"observed_data\": self.observed_data_to_xarray(),\r\n            }\r\n        )\r\n\r\n\r\ndef from_tfp(\r\n    posterior=None,\r\n    *,\r\n    var_names=None,\r\n    model_fn=None,\r\n    feed_dict=None,\r\n    posterior_predictive_samples=100,\r\n    posterior_predictive_size=1,\r\n    observed=None,\r\n    coords=None,\r\n    dims=None\r\n):\r\n    \"\"\"Convert tfp data into an InferenceData object.\"\"\"\r\n    return TfpConverter(\r\n        posterior=posterior,\r\n        var_names=var_names,\r\n        model_fn=model_fn,\r\n        feed_dict=feed_dict,\r\n        posterior_predictive_samples=posterior_predictive_samples,\r\n        posterior_predictive_size=posterior_predictive_size,\r\n        observed=observed,\r\n        coords=coords,\r\n        dims=dims,\r\n    ).to_inference_data()\r\n", "path": "arviz/data/io_tfp.py"}], "after_files": [{"content": "\"\"\"Tfp-specific conversion code.\"\"\"\r\nimport numpy as np\r\nimport xarray as xr\r\n\r\nfrom .inference_data import InferenceData\r\nfrom .base import dict_to_dataset, generate_dims_coords, make_attrs\r\nfrom .. import utils\r\n\r\n# pylint: disable=too-many-instance-attributes\r\nclass TfpConverter:\r\n    \"\"\"Encapsulate tfp specific logic.\"\"\"\r\n\r\n    def __init__(\r\n        self,\r\n        *,\r\n        posterior,\r\n        var_names=None,\r\n        model_fn=None,\r\n        feed_dict=None,\r\n        posterior_predictive_samples=100,\r\n        posterior_predictive_size=1,\r\n        chain_dim=None,\r\n        observed=None,\r\n        coords=None,\r\n        dims=None\r\n    ):\r\n\r\n        self.posterior = posterior\r\n\r\n        if var_names is None:\r\n            self.var_names = []\r\n            for i in range(0, len(posterior)):\r\n                self.var_names.append(\"var_{0}\".format(i))\r\n        else:\r\n            self.var_names = var_names\r\n\r\n        self.model_fn = model_fn\r\n        self.feed_dict = feed_dict\r\n        self.posterior_predictive_samples = posterior_predictive_samples\r\n        self.posterior_predictive_size = posterior_predictive_size\r\n        self.observed = observed\r\n        self.chain_dim = chain_dim\r\n        self.coords = coords\r\n        self.dims = dims\r\n\r\n        import tensorflow_probability as tfp\r\n        import tensorflow as tf\r\n        import tensorflow_probability.python.edward2 as ed\r\n\r\n        self.tfp = tfp\r\n        self.tf = tf  # pylint: disable=invalid-name\r\n        self.ed = ed  # pylint: disable=invalid-name\r\n\r\n        if int(self.tf.__version__[0]) > 1:\r\n            import tensorflow.compat.v1 as tf  # pylint: disable=import-error\r\n\r\n            tf.disable_v2_behavior()\r\n            self.tf = tf  # pylint: disable=invalid-name\r\n\r\n    def handle_chain_location(self, ary):\r\n        \"\"\"Move the axis corresponding to the chain to first position.\r\n\r\n        If there is only one chain which has no axis, add it.\r\n        \"\"\"\r\n        if self.chain_dim is None:\r\n            return utils.expand_dims(ary)\r\n        return ary.swapaxes(0, self.chain_dim)\r\n\r\n    def posterior_to_xarray(self):\r\n        \"\"\"Convert the posterior to an xarray dataset.\"\"\"\r\n        data = {}\r\n        for i, var_name in enumerate(self.var_names):\r\n            data[var_name] = self.handle_chain_location(self.posterior[i])\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n\r\n    def observed_data_to_xarray(self):\r\n        \"\"\"Convert observed data to xarray.\"\"\"\r\n        if self.observed is None:\r\n            return None\r\n\r\n        observed_data = {}\r\n        if isinstance(self.observed, self.tf.Tensor):\r\n            with self.tf.Session() as sess:\r\n                vals = sess.run(self.observed, feed_dict=self.feed_dict)\r\n        else:\r\n            vals = self.observed\r\n\r\n        if self.dims is None:\r\n            dims = {}\r\n        else:\r\n            dims = self.dims\r\n\r\n        name = \"obs\"\r\n        val_dims = dims.get(name)\r\n        vals = utils.one_de(vals)\r\n        val_dims, coords = generate_dims_coords(vals.shape, name, dims=val_dims, coords=self.coords)\r\n        # coords = {key: xr.IndexVariable((key,), data=coords[key]) for key in val_dims}\r\n\r\n        observed_data[name] = xr.DataArray(vals, dims=val_dims, coords=coords)\r\n        return xr.Dataset(data_vars=observed_data, attrs=make_attrs(library=self.tfp))\r\n\r\n    def _value_setter(self, variables):\r\n        def interceptor(rv_constructor, *rv_args, **rv_kwargs):\r\n            \"\"\"Replace prior on effects with empirical posterior mean from MCMC.\"\"\"\r\n            name = rv_kwargs.pop(\"name\")\r\n            if name in variables:\r\n                rv_kwargs[\"value\"] = variables[name]\r\n            return rv_constructor(*rv_args, **rv_kwargs)\r\n\r\n        return interceptor\r\n\r\n    def posterior_predictive_to_xarray(self):\r\n        \"\"\"Convert posterior_predictive samples to xarray.\"\"\"\r\n        if self.model_fn is None:\r\n            return None\r\n\r\n        posterior_preds = []\r\n        sample_size = self.posterior[0].shape[0]\r\n\r\n        for i in np.arange(0, sample_size, int(sample_size / self.posterior_predictive_samples)):\r\n            variables = {}\r\n            for var_i, var_name in enumerate(self.var_names):\r\n                variables[var_name] = self.posterior[var_i][i]\r\n\r\n            with self.ed.interception(self._value_setter(variables)):\r\n                if self.posterior_predictive_size > 1:\r\n                    posterior_preds.append(\r\n                        [self.model_fn() for _ in range(self.posterior_predictive_size)]\r\n                    )\r\n                else:\r\n                    posterior_preds.append(self.model_fn())\r\n\r\n        data = {}\r\n        with self.tf.Session() as sess:\r\n            data[\"obs\"] = self.handle_chain_location(\r\n                sess.run(posterior_preds, feed_dict=self.feed_dict)\r\n            )\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=self.dims)\r\n\r\n    def sample_stats_to_xarray(self):\r\n        \"\"\"Extract sample_stats from tfp trace.\"\"\"\r\n        if self.model_fn is None or self.observed is None:\r\n            return None\r\n\r\n        log_likelihood = []\r\n        sample_size = self.posterior[0].shape[0]\r\n\r\n        for i in range(sample_size):\r\n            variables = {}\r\n            for var_i, var_name in enumerate(self.var_names):\r\n                variables[var_name] = self.posterior[var_i][i]\r\n\r\n            with self.ed.interception(self._value_setter(variables)):\r\n                log_likelihood.append((self.model_fn().distribution.log_prob(self.observed)))\r\n\r\n        data = {}\r\n        if self.dims is not None:\r\n            coord_name = self.dims.get(\"obs\")\r\n        else:\r\n            coord_name = None\r\n        dims = {\"log_likelihood\": coord_name}\r\n\r\n        with self.tf.Session() as sess:\r\n            data[\"log_likelihood\"] = self.handle_chain_location(\r\n                sess.run(log_likelihood, feed_dict=self.feed_dict)\r\n            )\r\n        return dict_to_dataset(data, library=self.tfp, coords=self.coords, dims=dims)\r\n\r\n    def to_inference_data(self):\r\n        \"\"\"Convert all available data to an InferenceData object.\r\n\r\n        Note that if groups can not be created (i.e., there is no `trace`, so\r\n        the `posterior` and `sample_stats` can not be extracted), then the InferenceData\r\n        will not have those groups.\r\n        \"\"\"\r\n        return InferenceData(\r\n            **{\r\n                \"posterior\": self.posterior_to_xarray(),\r\n                \"sample_stats\": self.sample_stats_to_xarray(),\r\n                \"posterior_predictive\": self.posterior_predictive_to_xarray(),\r\n                \"observed_data\": self.observed_data_to_xarray(),\r\n            }\r\n        )\r\n\r\n\r\ndef from_tfp(\r\n    posterior=None,\r\n    *,\r\n    var_names=None,\r\n    model_fn=None,\r\n    feed_dict=None,\r\n    posterior_predictive_samples=100,\r\n    posterior_predictive_size=1,\r\n    chain_dim=None,\r\n    observed=None,\r\n    coords=None,\r\n    dims=None\r\n):\r\n    \"\"\"Convert tfp data into an InferenceData object.\"\"\"\r\n    return TfpConverter(\r\n        posterior=posterior,\r\n        var_names=var_names,\r\n        model_fn=model_fn,\r\n        feed_dict=feed_dict,\r\n        posterior_predictive_samples=posterior_predictive_samples,\r\n        posterior_predictive_size=posterior_predictive_size,\r\n        chain_dim=chain_dim,\r\n        observed=observed,\r\n        coords=coords,\r\n        dims=dims,\r\n    ).to_inference_data()\r\n", "path": "arviz/data/io_tfp.py"}]}