Problem with Boolean masks indexing by jax jit

[alisheikholeslam]

I have written a code as:

from functools import partial
import jax
import jax.numpy as jnp
from jax import jit, vmap
import numpy as np


@partial(jit, static_argnums=(1,))
def moving_window(a, size: int):
    starts = jnp.arange(len(a) - size + 1)
    return vmap(lambda start: jax.lax.dynamic_slice(a, (start,), (size,)))(starts)

@jit
def jax_comp(a, b):
    shape, size = a.shape[0] + 1, a.shape[0] + b.shape[0]
    ind_0 = moving_window(jnp.arange(size), b.shape[0])
    ind_1 = moving_window(jnp.arange(b.shape[0], size + shape - 1), a.shape[0])
    ind_1 = jnp.remainder(ind_1, size)                    # ind_1 = ind_1 % size
    arr = jnp.zeros((shape, size), dtype=jnp.float64)
    arange_ = jnp.arange(shape)[:, None]
    arr = arr.at[arange_, ind_0].set(b)
    arr = arr.at[arange_, jnp.sort(ind_1)].set(jnp.broadcast_to(a, (shape, a.shape[0])))
    return arr

a = np.random.rand(1000)
b = np.array([5, 6, 10])

I tried to write this code in to other forms using jax jit, but it stuck when reaching Boolean masks, which, I think, could be handled by jnp.where:

@jit
def jax_(a, b):
    shape, size = a.shape[0] + 1, a.shape[0] + b.shape[0]
    ind = moving_window(jnp.arange(size), b.shape[0])
    arr = jnp.zeros((shape, size), dtype=jnp.float64)
    arr = arr.at[jnp.arange(shape)[:, None], ind].set(b)
    # arr = arr.at[arr == 0].set(jnp.broadcast_to(a, (shape, a.shape[0])))
    arr = jnp.where(arr == 0, jnp.broadcast_to(a, (shape, a.shape[0])), arr)
    return arr

@jit
def jax_2(a, b):
    shape, size = a.shape[0] + 1, a.shape[0] + b.shape[0]
    res = jnp.empty([shape, size])
    frame = jnp.reshape(jnp.arange(shape * size), (shape, size))
    diag_mask = (frame % (res.shape[1] + 1)) < (b.shape[0])
    # res = res.at[diag_mask].set(jnp.tile(b, res.shape[0]))
    # res = res.at[~diag_mask].set(jnp.tile(a, res.shape[0]))
    res = jnp.where(diag_mask, jnp.tile(b, res.shape[0]), jnp.tile(a, res.shape[0]))
    return res

which get such the following error:

ValueError: Incompatible shapes for broadcasting: ((1001, 1005), (1001, 1000), (1001, 1005))

What is the problem and how can I pass it?

[jakevdp]

The reason you're getting this error is because the three-argument jnp.where function requires that the arguments be broadcast-compatible, and your arguments are not broadcast-compatible (they have differing shapes).

It sounds like what you're after is to have a JIT-compatible version of something like this:

import jax
import jax.numpy as jnp

x = jnp.arange(10)
y = jnp.ones(5)

def f1(x, y):
  return x.at[x >= 5].set(y)

print(f1(x, y))
# [0 1 2 3 4 1 1 1 1 1]

When you try to JIT-compile this, it errors:

# print(jax.jit(f1)(x, y))
# NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[10])

The reason for this is that there's no way to know at compile-time how many elements of x will satisfy the condition, and so there is no way to know whether the code is valid.

The alternative you tried is roughly this:

def f2(x, y):
  return jnp.where(x >= 5, y, x)

This is invalid because x is of size 10 and y is of size 5, so they are not broadcast compatible, as is required by the semantics of the where function.

Instead you need to do something like this:

def f3(x, y):
  idx = jnp.where(x >= 5, size=len(y))
  return x.at[idx].set(y)

print(f3(x, y))
# [0 1 2 3 4 1 1 1 1 1]

print(jax.jit(f3)(x, y))
# [0 1 2 3 4 1 1 1 1 1]

Notice that this involves specifying at compile time that you expect there to be 5 indices where the condition is True. This information allows JAX to correctly JIT-compile the code.

[alisheikholeslam]

I have wrote the second code as:

@jit
def jax_evader(a, b):
    shape, size = a.shape[0] + 1, a.shape[0] + b.shape[0]
    res = jnp.empty([shape, size])
    frame = jnp.reshape(jnp.arange(shape * size), (shape, size))
    diag_mask = (frame % (res.shape[1] + 1)) < (b.shape[0])
    res_0 = jnp.tile(b, res.shape[0])
    res_1 = jnp.tile(a, res.shape[0])
    
    idx = jnp.where(diag_mask, size=res_0.size)
    idx_v = jnp.where(~diag_mask, size=res_1.size)
    res = res.at[idx].set(res_0)
    return res.at[idx_v].set(res_1)

    # return jnp.where(diag_mask, res_0, res_1)    < -- ?????

Can I write these two pair of jnp.where and assignment in a compacted way as I have shown on the code? Any recommendation is appreciated. (unrelated question: how to make codes colorful as you have wrote. :))