Is there a jittable way to check for non zero values in an array?

[bsaoptima]

Hi there,

I have a function that extracts non-zero diagonals from an input square matrix. It would be better if I wrap it with @jax.jit but I am getting stuck at a step where I actually have to check if the diagonals contain any non-zero values as I have conditional logic based on that condition.

def to_sparse(other):
    diagonals = []
    offsets = []

    n = other.shape[0]

    if not isinstance(other, Array):
        other = jnp.asarray(other.array)

    for offset in range(-n + 1, n):
        diagonal = jnp.diagonal(other, offset=offset)
        if jnp.any(diagonal != 0):
            diag = jnp.zeros((n,))
            start = max(0, offset)
            end = min(n, n + offset)
            diagonals.append(diag.at[start:end].set(diagonal))
            offsets.append(offset)

    diagonals = jnp.array(diagonals)
    offsets = tuple(offsets)

    return NewObject(diagonals, offsets)

So the conditional logic based on jnp.any fails, I tried an implementation with jax.lax.cond but since I append stuff to a list, it fails. I'm going around in circle thus why I am posting here.

Anyone could help me out with this? Thank you in advance 😃

[jakevdp]

No, there is no way to use an if statement with traced values. See JAX sharp bits: control flow for some general discussion of this. Briefly, the issue is that Python if statements must be evaluated at compile time, but the values in the arrays will not be known until runtime.

The best alternative is to use jax.lax.cond, which is essentially a runtime if-statement. However, this will require both branches of your cond to have the same shape, so appending in one and not the other is not possible. Typical strategies in this case would be to use a static array to store offsets rather than appending them dynamically.

Hope that helps!

[bsaoptima]

Oh okay thank you very much! If offsets was a tuple, would the appending work? I found a way to split this function in 2 parts (one jitted) where I do use offsets as a static argument but I would rather have a fully jitted function.

[jakevdp]

You can append to the list in cond as long as both branches of the cond append. For example this will work:

jax.lax.cond(condition,
    lambda x, y: x.append(y),
    lambda x, y: x.append(0),
    x, y)

This will not work:

jax.lax.cond(condition,
    lambda x, y: x.append(y),
    lambda x, y: x,
    x, y)

The reason for this is that condition is not known until runtime, so if the cond returns a different number of results based on the condition, there's no way to represent the output in JAX's abstract IR prior to the runtime values being known.

If you want to jit-compile your full algorithm, you'll have to change some of the details: for example using placeholders or pad values rather than conditionally appending values.

[bsaoptima]

This makes sense, I am trying to use placeholders, just have to find the right formula, thanks for the help!