Pm.sample result does not work with az.plot_trace or other az functions

Not sure how that happened as your code shouldn’t have produced such a result. Somehow your idata object is wrapped in a tuple (notice the parentheses around the result). When I run a simplified version of you model, this is what I get:

In [3]: idata
Out[3]: 
Inference data with groups:
	> posterior
	> log_likelihood
	> sample_stats
	> observed_data