Where is the source code for kl_div?

I’m looking @ https://pytorch.org/docs/stable/_modules/torch/nn/functional.html

to see torch.nn.kl_div but it is just a class that wraps a function called torch.kl_div
I can’t find that original function. I basically made my own function and it spits out different results from what Pytorch built-in is spitting so I’m wondering how does it look like.

For completition, here’s my code:

def my_kl(predicted, target):
    return -(target * t.log(predicted.clamp_min(1e-7))).sum(dim=1).mean() - \
           -1*(target.clamp(min=1e-7) * t.log(target.clamp(min=1e-7))).sum(dim=1).mean()

Which I believe is spitting out correct results LOL

kl_div and the CPU backward are in aten/src/ATen/native/Loss.cpp, the cuda backward is in aten/src/ATen/native/cuda/Loss.cu.
I once tried to write this up more generally:
https://lernapparat.de/selective-excursion-into-pytorch-internals/

Best regards

Thomas

2 Likes