# How to use numpy for other general library to rewrite nn.KLDiv?

In my project, I used nn.KLDiv(A.log(), B) to calculate the KL divergence between A and B.
But now I have to rewrite this part as numpy or other general python libraries in order to some model conversion.

I did a lot of search on google, and tried the following implementations:

``````(1) First method
def KL(a, b):
a = np.asarray(a, dtype=np.float)
b = np.asarray(b, dtype=np.float)
return np.sum(np.where(a != 0, a * np.log(a / b), 0))

(2) Second method
scipy.special.kldiv
Neither of them show the similar output with the nn.KLDivLoss method in pytorch.
``````

Could someone provide some guidance how to implement KL divergence with numpy to provide similar output with pytorch? Thanks.

Using the formula from the docs, I get the same results (up to floating point precision):

``````def KL(a, b):
return b * (np.log(b) - a)

kl_loss = nn.KLDivLoss(reduction="batchmean")
# input should be a distribution in the log space
input = F.log_softmax(torch.randn(3, 5, requires_grad=True), dim=1)
# Sample a batch of distributions. Usually this would come from the dataset
target = F.softmax(torch.rand(3, 5), dim=1)
ref = kl_loss(input, target)

# numpy reference
out = KL(input.detach().numpy(), target.numpy())
out = np.sum(out) / out.shape[0]

print((ref - out).abs().max())