Filter an array based on another boolean array (JIT compiled)

[valentinmace]

I think my request might be inherently impossible to achieve but I still want to ask. This is a simplified version of my case:

I have an array of points, say points = jnp.array([[1,1], [2,2], ... [9,9]]) and I compute a boolean array to filter my points array, say filter = jnp.array([True, False, ..., True])

I want create a new array such that only points which have a True corresponding value in filter appear in this new array.

In standard python this is trivial, you can do filtered_points = points[filter, :]

However I cannot manage to find a solution in compiled jax with jit since the shape of filtered_points will vary depending on the boolean values in filter.

The reason why I don't want to use jnp.where to keep fixed sized arrays and for example replace False points by nans in filtered_points is that I'm then computing euclidian distances of filtered_points against itself. Since the computation of these distances has a complexity of O(N^2) I want filtered_points to be as short as possible and replacing some elements by nans won't speed up this computation.

[jakevdp]

No, unfortunately jax.jit requires statically-shaped arrays, so there is no way to do what you're asking within a JIT-compiled function. The standard workaround is to re-express your computation in terms of statically shaped arrays, and I find in practice this can often be done with some thought. Another possiblity is to split your computation such that the dynamic arrays are generated outside of a JIT context, and the computation on those arrays are then done within JIT.

If you have a specific application in mind that you're having trouble re-expressing this way, feel free to open a discussion with more detail.

[valentinmace]

Thanks for you responsiveness.

I can do what I want using statically shaped arrays. My objective was to accelerate computation of euclidian distances which is the next step (expensive in terms of memory and time) in my program by providing shorter arrays of points.

In practice, filtered_points often contains a lot less points than the original points array but there is no way to predict how many. In standard python, the filtering helps to drastically speed up the process and I was hopping to find a way in compiled jax too.

[YouJiacheng]

Hey!
For acceleration purpose, I think there is a solution.

1. do something like c++ std::erase, move all invalid value to the tail of array while preserving shape.
2. reshape the array into chunks, i.e. (n, ...) -> (num_chunks, n // num_chunks, ...).
3. use lax.while_loop (or lax.scan with lax.cond if reverse-mode autodiff is needed) to only compute for valid chunks.

But it doesn't help if the computation isn't much more expensive than memory access.

import jax
from jax import numpy as jnp
from jax import lax

def jax_filter(conds, xs):
    """
        returns:
            result array
            number of valid elements
    """
    res = jnp.full_like(xs, jnp.nan)
    def body_fun(carry, v):
        res, idx = carry
        cond, x = v
        # there is no need for a jnp.where or lax.cond
        # since only [0, last_idx) is valid
        return (res.at[idx].set(x), idx + cond), None
    return lax.scan(body_fun, (res, 0), (conds, xs))[0]

def valid_chunks_vmap(f, num_chunks: int):
    fs = jax.vmap(f)
    def valid_chunks_f(xs, num_valid_chunks):
        def body_fun(idx, x_chunk):
            def no_op(x):
                s = jax.eval_shape(fs, x)
                return jnp.full(s.shape, jnp.nan, s.dtype)
            return idx + 1, lax.cond(idx < num_valid_chunks, fs, no_op, x_chunk)
    
        y_chunks = lax.scan(body_fun, 0, jnp.reshape(xs, (num_chunks, -1, *xs.shape[1:])))[1]
        return jnp.reshape(y_chunks, (-1, *y_chunks.shape[2:]))
    return valid_chunks_f

def f(x, y):
    return jnp.sum((x - y) ** 2)

@jax.jit
def filtered_distances(conds, xs):
    xs, n_valid = jax_filter(conds, xs)
    num_chunks = 16
    assert xs.shape[0] % num_chunks == 0
    chunk_size = xs.shape[0] // num_chunks
    num_valid_chunks = (n_valid + chunk_size - 1) // chunk_size
    def inner(x):
        return valid_chunks_vmap(lambda z: f(x, z), num_chunks)(xs, num_valid_chunks)
    return valid_chunks_vmap(inner, num_chunks)(xs, num_valid_chunks)

xs = jnp.arange(256, dtype=jnp.float32)
conds = (jnp.arange(256) % 2).astype(jnp.bool_)
print(filtered_distances(conds, xs))

[valentinmace]

Thanks a lot for your answer, I'll give it a try !

[jakevdp]

I suspect it will be faster to replace jax_filter with something like this, particularly on accelerators where scan is slow:

def jax_filter(conds, xs):
  return xs[jnp.where(conds, size=len(xs))], conds.sum()

[YouJiacheng]

Thanks @jakevdp , I forget jnp.where(or jnp.nonzero) has such usage!

[YouJiacheng]

@jakevdp I wonder if following code would faster than jnp.where version?
jnp.where use cumsum, bincount(.at[x].add(1)) and cumsum.

def jax_filter(conds, xs):
    cumsum = jnp.cumsum(conds)
    return jnp.zeros_like(xs).at[cumsum - 1].add(jnp.where(conds, xs, 0)), cumsum[-1]

[YouJiacheng]

oops, as @jakevdp pointed out, scan is extremely slow.

import jax
from jax import numpy as jnp
from jax import lax

@jax.jit
@jax.vmap
def ff_1(conds, xs):
    cumsum = jnp.cumsum(conds)
    return jnp.zeros_like(xs).at[cumsum - 1].add(jnp.where(conds, xs, 0)), cumsum[-1]

@jax.jit
@jax.vmap
def ff_2(conds, xs):
    cumsum = jnp.cumsum(conds)
    intervals = jnp.zeros_like(cumsum).at[cumsum].add(1)
    idx = jnp.cumsum(intervals)
    return xs[idx], cumsum[-1]

@jax.jit
@jax.vmap
def ff_3(conds, xs):
    return xs[jnp.where(conds, size=len(xs))], conds.sum()

@jax.jit
@jax.vmap
def ff_4(conds, xs):
    res = jnp.full_like(xs, jnp.nan)
    def body_fun(carry, v):
        res, idx = carry
        cond, x = v
        return (res.at[idx].set(x), idx + cond), None
    return lax.scan(body_fun, (res, 0), (conds, xs))[0]

bsz = 1024
n = 65536
xs = jnp.broadcast_to(jnp.arange(n, dtype=jnp.float32), (bsz, n))
conds = jnp.broadcast_to(lax.convert_element_type(jnp.arange(n) % 2, jnp.bool_), (bsz, n))

def timer(f):
    from time import time
    f() # warmup
    t = time()
    for _ in range(3):
        f()
    print((time() - t) / 3)

timer(lambda: jax.block_until_ready(ff_1(conds, xs)))
timer(lambda: jax.block_until_ready(ff_2(conds, xs)))
timer(lambda: jax.block_until_ready(ff_3(conds, xs)))
timer(lambda: jax.block_until_ready(ff_4(conds, xs)))

0.0033173561096191406
0.0063737233479817705
0.006862243016560872
7.318466663360596

8x larger n for the first three:

0.024444182713826496
0.048021952311197914
0.04943092664082845