Strange it does not work if the function is outside the context. Do you think it is some sort of bug?
I implemented the function as you suggested and it seems that works (see updated notebook). Note that I removed broadcast_shape= because it seems it does not change the results. What do you mean with
If I run ppc = pm.sample_posterior_predictive(trace, samples=250, size=100) I get an (250x100) array, which I think is correct.