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.

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.

[junpenglao]

I am intend to say not supporting, I haven’t really seen any valid use case.

[OriolAbril]

Should we support multiple chain dims? We could stack them before conversion to dataset