result = safe_softmax(jnp.array([1.0, 2.0, -1.0]))
chex.block_until_chexify_assertions_complete()

        __
        __