Loss Fn - Kullback-Leibler divergence (relative entropy) loss

[cksac]

ref: https://optax.readthedocs.io/en/latest/api/losses.html#optax.kl_divergence