@ericmjl has great content all over YouTube. However, it is this blog post that has proved most helpful to me while wrestling with shape dimensions in NumPyro : "Reasoning about Shapes and Probability Distributions". So good! https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/