How to do value assertion gracefully in JAX?

[HeavyCrab]

I find it hard to use assertion in JAX.
For static assertions, we can do it gracefully with chex.
However it seems hard to use value assertion with chex.chexify, for example, it should only wrapping jitted function and should not be nested.
So I use this workaround to do my own value assertions. It should work in both jitted and unjitted functions.

def jax_assert(cond, msg: str):
    def _raise(): raise ValueError(msg)
    def raise_branch():
        jax.debug.callback(_raise)
        return 0
    return jax.lax.cond(cond, lambda: 0, raise_branch)

And the assertion can be easily disabled by:

jax_assert = lambda cond, msg: None

However, according to the document of jax.lax.cond:

In contrast with jax.lax.select(), using cond indicates that only one of the two branches is executed (up to compiler rewrites and optimizations). However, when transformed with vmap() to operate over a batch of predicates, cond is converted to select(). Both branches will be traced in all cases (see Key concepts: tracing for a discussion of JAX’s tracing model).

This workaround may fail when using with jax.vmap, since each branch will be traced, and therefore the error will be thrown automatically even with a True cond. I wonder if there is a simple way to do value assertion in JAX. How do you guys do value assertions in JAX?

[jakevdp]

Unfortunately, there's no supported way to do runtime value assertions in JAX. One main reason for this is that some accelerators supported by JAX literally have no path for raising errors at runtime.

You can do various imperfect workarounds using callbacks, but those are not recommended. Better would be to modify your programming style to avoid the need for runtime value assertions. It takes some getting used to, but we find it works pretty well in practice.

Another option is to use the jax.experimental.checkify module, which allows you to propagate value-based errors using a functional approach supported by JAX. It's still somewhat rough around the edges, but there's some ongoing work to make the checkify-style approach more seamless and convenient.

[patrick-kidger]

Actually btw Jake, I thought I'd mention here that last time I checked (about a year ago), eqx.error_if actually does support TPUs. The trick is to chain a second callback after the first, so that if the error in the first is ignored then the second callback spins in an infinite loop instead:

https://github.com/patrick-kidger/equinox/blob/d3d4eb9758a017b1caddbc5219fd09379e171681/equinox/_errors.py#L42

I know you're not a fan of the callback-based approach, but I thought I'd at least highlight that I think that particular problem is solved.

[jakevdp]

Hi Patrick – thanks for weighing in. I'm aware of error_if and the tricks that equinox plays to work around the limitations of the TPU and other accelerators. However, solutions like this — which depend on potentially ephemeral implementation details rather than guaranteed API contracts — are not something I'd recommend for JAX users.

[patrick-kidger]

Yup, I'm explicitly not here looking to recommend it in this thread / I know we've both mentioned these options in these kinds of threads before.

Rather, I just wanted to let you know that I think specifically errors-on-TPU is solved, as I read your previous comment to mean that detail might not be something you'd seen before.

(In fact AFAIK this fallback option means that we don't actually rely on any noncontractual details -- actually raising a Python error is certainly a nice-to-have, but preventing further computation is really the important part -- but anyway not super important.)

:)