在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])来加快速度，但这似乎矢量化程度仍然很低。我似乎找不到任何可以做得更好的内置方法。