import chex

# chexify goes OUTSIDE jit

@chex.chexify
@jax.jit
def safe_softmax(x):
    probs = jax.nn.softmax(x)
    
    chex.assert_equal(
        jnp.all(probs >= 0), 
        True, 
        custom_message="Negative probability detected!"
    )
    
    return probs

        __
        __