jax

c3c23932 - kl_div: fix incorrect formula in doc