You just need to reshape your type_idx into shape (3,1) instead of (3,):
type_idx = np.array((0,0,1)).reshape(3,1)
That being said, in both your models you get a ton of divergences, potentially leading to biased estimates (even though it does capture your tau_true well). You might want to look at reformulating your model a bit. This is probably down to the fact that your noise term is infinitesmal, on the order of 10^-11, so your observations aren’t really Normal, though this could be a quirk of your toy data here.