I'm not sure if this is supposed to be a matmul or vecdot. I think the later. If someone is familiar with the distribution please weigh in.

CC @bwengals or perhaps @jessegrabowski (in case we can simplify this with PyTensor magic)