I actually tried to do the example with a batch dimension and it didn’t work, I imagine because of exactly the issue you raise.
I really wanted something to clearly demonstrate the point though, because to this seems to be the important thing to grok: given priors from the loc-scale family, a nested hierarchical model can be pulled apart into a series of sums and offsets. This, combined with fancy indexing, makes it easy to work with ragged data structures.