jax

Utility functions for JAX models.

Functions

jax_entropy	Shannon entropy H(p) computed in jax along the last dim; 0*log(0) treated as 0.
jax_fresh_prng_key	Draw a fresh PRNG key seeded from OS entropy.
jax_intersection_probability	Intersection probability of a probability interval, per [WCM+24] Section 3.4.