It would be great if there were a way to handle label switching in traces. I’ve come up against this with a latent space model where the dimensions are different across traces. In this case I’d been worried about rotation and shearing, but it appears that the dimensions are consistent across chains, they’re just mixed and matched. I’ve seen other discussion on the Discourse where people have had similar issues, and I have some ideas for how to handle it.
By exploring the traces themselves you can figure out which dimensions should align, but I haven’t figured out how to change the underlying data in the trace object. For example, the following does not work. Here, I use a 3000 sample trace for a 3 dimension latent variable (size 500x3). I try to overwrite the data in the trace object such that the sampled values for the 0th and 1st dimensions are swapped. When I then try to retrieve those values, they are unaltered.
trace2 = copy.copy(trace) trace2.z[3000:5999,:,0]=trace.z[3000:5999,:,1] trace2.z[3000:5999,:,1]=trace.z[3000:5999,:,0]
Apologies if there’s already an easy way to do this: dir(trace) lists a method to remove values, but not to change values. This would be one possible workaround, because after switching the dimensions in the trace itself, the diagnostics would all work properly.
Another possibility would be to allow pm.summary() and other diagnostics to take an array instead of only accepting a multitrace object.
For example, in diagnostics.py on line 172, the gelmin_rubin function just converts it into a numpy array.
x = np.array(mtrace.get_values(var, combine=False))
The function could easily just take x directly, perhaps with a warning to make sure you know what you’re doing if you don’t use a multitrace. For my case, I just copied the functions from the pymc3 source and used numpy arrays to calculate the diagnostics, but it seems like a small change that would make pymc3 more useful functional for models with latent variables.
Using arviz is another solutions here, but I had trouble figuring out the dimensions for the data structures, and how to keep chains separate. So the easiest solution might just to add a few lines to the documentation of the arviz package (so I’ll also tag @RavinKumar )