В Pytorch у меня есть три тензора A, mu и sigma, все с формой B x N x M.
Я хочу новый тензор D формы B x N x M x B,где каждая запись - это нормальная плотность записи A относительно пары параметров в mu и sigma.Более конкретно:
D[i,j,k,l] = univariate_normal_density(A[i,j,k] | mu[l,j,k], sigma[l,j,k])
Другими словами, каждая запись mu и sigma представляет собой среднее значение и дисперсию распределения.Кстати, не имеет значения, содержит ли D вероятности или логарифмические вероятности.Также обратите внимание, что можно развернуть второе и третье измерения, я думаю, что я чего-то там упускаю.
Какой самый эффективный способ сделать это?
Наиболее простой способ выглядит так:
Создание B*N*M объектов типа torch.distributions.normal, так что N[l,j,k]=N(mu[l,j,k],sigma[l,j,k])
Для каждогооцените log_prob на A как D[i,j,k,l] = N[l,j,k].log_prob(A[i,j,k])
Я могу ускорить это с помощью: D[:,j,k,l] = N[l,j,k].log_prob(A[:,j,k]), но это все еще кажется недостаточным.Я не могу найти какие-либо встроенные методы, которые могли бы работать лучше.