How to display the trace of the diagonal element of a matrice RV

Hi,
I have a model that involve a covariance matrix named CovMatrix. When I want to plot the trace of it using pm.traceplot(trace, var_names=['CovMatrix'], I have of course one trace for each element in my matrix. How can I display the trace of the diagonal element only? I think that I should use the coords argument but I am not sure how to use it for my particular case.
Thanks for the help!

Hi Timothé!
Usually, when I need to use ArviZ’s coords arguments, I transform my trace into an xarray (ArviZ uses that in the backend): az.from_pymc3(my_trace), and if you look at az.from_pymc3(my_trace).posterior you’ll see the different dimensions, that’s really helpful to do trial and error about which coords you need.

Alternatively, I think you can also directly isolate the diagonal of you covariance matrix in the model. If it’s a Cholesky decomposition for instance:

cov = pm.math.dot(chol, chol.T)
var_cluster = pm.Deterministic("var_cluster", tt.diag(cov))