All of the model code looks fine. You will just need to familiarize yourself with the new Arviz InferenceData object that is returned from sampling. See here for documentation about idata, or here for a discussion with a user who had a similar problem.
Amusingly (in the context of this thread), the Arviz idata docs reference PyMC3, but there will not be any differences with respect to working with idata.