jax_entropy¶

probly.utils.jax.jax_entropy(p: Array) → Array[source]¶

Shannon entropy H(p) computed in jax along the last dim; 0*log(0) treated as 0.

The logarithm is fed with the zeros replaced by ones instead of masking its result, so the gradient stays finite for probability vectors that contain exact zeros.

Parameters:: p – Probabilities to compute entropy of.
Returns:: Entropy of probabilities p.