Checkify in vmaped lax.cond/lax.select.

[SobhanMP]

I've been trying to debug a code with Checkify and it is great but what would be the "idiomatic" way of dealing with a code like this

import os
from typing import Any
os.environ['CUDA_VISIBLE_DEVICES'] = ''
import jax
import jax.numpy as jnp
from jax import vmap, jit, lax
from jax.experimental import checkify
def f(x, y):
    def f_true(x):
        checkify.check(y == 0, 'y is {} not 0', y)
        return x * 2
    def f_false(x):
        checkify.check(y != 0, 'y is not zero but {}', y)
        return x * 3
    return lax.cond(x == 0, f_true, f_false, x, y)
f = checkify.checkify(jit(vmap(f)), errors=checkify.all_checks)
err, x = f(jnp.array([0]), jnp.array([0]))
err.throw()

Since the lax.cond gets converted into a lax.select, both branches are executed. As a result both checks are executed and one of them fails.

Is there a way of fixing this inside jax? Is the solution just making the checks like

checkify.check(jnp.logical_or(x != 0, y == 0), 'y is {} not 0', y)

?

[patrick-kidger]

I think this will take a bit more work to tackle in general in JAX. The checkify was presumably there for good reason: simply disabling it might mean that e.g. that branch results in an infinite loop.

So in general I think you probably want to do something like

def f_true(x, _):
    ...

def f_false(_, x):
    ...

pred = x == 0
safe_true_x = jnp.where(pred, x, some_safe_value)
safe_false_x = jnp.where(pred, some_other_safe_value, x)
lax.cond(pred, f_true, f_false, safe_true_x, safe_false_x)

where some_safe_value and some_other_safe_value are something you know will pass any checkifys in f_true and f_false, not get caught in any infinite loops, etc.

Side note, you may like error_if as an alternative to checkify.

[SobhanMP]

Interesting, I hadn't considered checks that would block an infinite loop, but what you say makes sense.

Do you think it's possible to get the checkify error only from the selected branches?

It would have been nice if there were an EQX_ON_ERROR to do nothing; the nan strategies only work with floats. I also like the fmt string in Checkify.

p.s. thanks for the fast answer!

[patrick-kidger]

Do you think it's possible to get the checkify error only from the selected branches?

Not that I know of.

It would have been nice if there were an EQX_ON_ERROR to do nothing; the nan strategies only work with floats.

I considered this, but thought the NaN strategy would be a little safer. Other dtypes will have some other value used, e.g. the maximum value for an integer.

I also like the fmt string in Checkify.

Haha, we were just talking about adding this in patrick-kidger/equinox#644

[SobhanMP]

I considered this, but thought the NaN strategy would be a little safer. Other dtypes will have some other value used, e.g. the maximum value for an integer.

That's neat; with haiku in maintenance mode, I guess it's time to move to equinox! I guess I was thinking of pandas and their dreadful nans.