jax

0501d3d7 - Merge pull request #21223 from jakevdp:kl-div-doc