It should not really matter. Some plots and other functions will need to have the original chain, draw shape, but you can always unstack. I personally prefer not stacking unless necessary, but as we say in catalan “cada mestre té el seu llibre” (literal translation “every teacher has their own book”)
There are several issues in the sample code.
This is sorting treatment and mu independently! You therefore have no guarantees that the 1-1 relation between treatment and their corresponding mu is maintained. You can use np.argsort and handle broadcasting, :… but again xarray can help you with that. Take a look for example at code cell 21 of https://docs.pymc.io/notebooks/multilevel_modeling.html.
This is not sorted at all, unlike the plot above, so you should not expect those to be comparable at all, ever. You are also plotting the HDI for mu, not for y, so you are plotting the hdi of the mean of your predictions, not of the predictions themselves. These are very different things. See for example code cell 16 in https://docs.pymc.io/notebooks/multilevel_modeling.html.