> Personal Example: I needed to sample a bunch of datapoints, make distributions out of them, sample, and then compute the density of each sample across distributions. Doing this with pytorch was rather slow, I was probably doing something wrong with vectorization and broadcasting, but I didn't have the time to figure it out.

You probably were not doing anything wrong. I spent a lot of time trying to be clever in order to parallelize things like this and it just wasn't possible without doing CUDA extensions. But it is now! PyTorch now has vmap through functorch and it works.