Cross-entropy computations trigger checkify NaN error

[hr0nix]

JAX recently added a transformation that allows to validate user-defined assertions, as well as check for some standard issues, such as division by zero or NaNs arising in computations. Unfortunately, some of this checks are currently violated by the distrax code. For example, this test currently passes:

def test_distrax_ce_nan():
    def func_to_check(logits, labels):
        return distrax.Categorical(probs=labels).cross_entropy(distrax.Categorical(logits=logits))
    checked_func = checkify.checkify(func_to_check, errors=checkify.nan_checks)
    err, val = checked_func(logits=jnp.log(jnp.asarray([0.5, 0.5])), labels=jnp.asarray([1.0, 0.0]))
    with pytest.raises(ValueError):
        err.throw()

The problem lies in the way cross-entropy is computed when some of the probabilities are zero: the current implementation computes zero times negative infinity, which results in NaNs, and then filters the NaNs using jnp.where. Unfortunately, while correct, this implementation triggers a checkify error because of these intermediate NaNs. This becomes problematic when one wants to enable NaN checks at the top-level.

A list of possible solutions that come to mind:

• Changing cross-entropy computation to avoid computing NaNs, for instance by replacing negative infinities in logits with finite values (the result will not be affected anyway because of jnp.where).
• Adding a new transformation to JAX that can be used to declare that a particular function triggers a check but that's OK (a somewhat related comment here: BUG: checkify of a vmap-ed while loop fails with an obscure error jax-ml/jax#11905 (comment))

Thanks!