Pytorch: Fast create score tensor method
I am new to Pytorch and am looking for a quick get score function. That, given a bunch of samples and a distribution, outputs a tensor consisting of the corresponding score for each individual sample. For instance, consider the following code:
norm = torch.distributions.multivariate_normal.MultivariateNormal(torch.zeros(2),torch.eye(2))
samples = norm.sample((1000,))
samples.requires_grad_(True)

Using samples I would like to create a score tensor of size [1000,2] where the ith component score[i] is the gradient of log p(samples[i]), where p is the density of the given distribution. The method I have come up with is the following:
def get_score(samples,distribution):
    log_probs = distribution.log_prob(samples)
    for i in range(log_probs.size()[0]):
        log_probs[i].backward(retain_graph = True)

The resulting score tensor is then samples.grad. The issue is that my method is quite slow for larger samples (e.g. for a sample of size [50000,2] it takes about 25-30 seconds on my CPU). Is this as fast as it can get?
The only alternative I can think of is to hard-code the score function for each distribution I will use, this doesn't seem like a good solution!

From experimentation, for 50000 samples, the following is about 50% quicker:
for i in range(50000):
    sample = norm.sample((1,))
    sample.requires_grad_(True)
    log_prob = norm.log_prob(a)
    log_prob.backward()

This indicates that there should be a better way!

I'm assuming that log_probs is stored as a pytorch tensor.
You can take advantage of the linearity of differentiation to calculate the derivative for all samples at once: log_probs.sum().backward(retain_graph = True)
At least with GPU acceleration this will be a lot faster.
If log_probs is not a tensor but a list of scalars (represented as pytorch tensors of rank 0), you can use log_probs = torch.stack(log_probs) first.
