distrax (pypi): Distrax: Probability distributions in JAX.
Keywords: jax|probability|distribution|python|machine|learning