Statistical Rethinking 2nd Ed, Ch4, Exercise 4H1

Your reasoning is correct, your first block of code computes mu for each of the new observations, not the posterior predictive distribution.

The easiest way to handle this is to use a pm.MutableData container to hold the training data, then use pm.set_data to swap in the out-of-sample data. There’s an example of this here.