Equivalent of jnp.where for arguments of different lengths

[SNMS95]

I am trying to optimize a piece of code by trying to JIT it. The code is as follows:

mask = [1, 1, 0, 1, 1, 0]
x = [10, 11, 12, 13] # Cannot be the same shape as mask
y = [1.0]

output = []
xcount = 0
for i in range(len(mask)):
    if mask[i] == 1:
        output.append(x[xcount])
        xcount+=1
    else:
        output.append(y[0])

What I am trying to achieve is to select elements from x or y based on the mask. If x was broadcastable to the shape of mask, I could have used jnp.where(mask, x, y) and could JIT this easily. However, due to some constraints, the length of x needs to be equal to the number of 1s present in mask. Is there a way out of using the manual for loop? I know that I can't use the fori_loop since the body function depends on the value of the mask. It might of interest to know that the mask variable is static while x can change, hence the need for JITing.

[jakevdp]

It looks like what you're after are the exact semantics of jax.numpy.place:

import jax.numpy as jnp
output = jnp.place(
    jnp.full(len(mask), y[0]),
    jnp.array(mask, dtype=bool),
    jnp.array(x), inplace=False)

This was only recently implemented in JAX (it will be part of the 0.4.15 release) but until that's available the implementation is pretty simple: https://github.com/google/jax/blob/d7940ee9a11d27fa5dc745e23e6f82aea25ac525/jax/_src/numpy/lax_numpy.py#L4985-L4987

For your particular case, you could use something like this:

mask = jnp.array(mask, dtype=bool)
x = jnp.array(x)
updates = jnp.zeros_like(x, shape=len(mask)).at[:len(x)].set(x)
indices = jnp.where(mask, size=mask.size, fill_value=mask.size)[0]
output = jnp.full(len(mask), y[0]).at[indices].set(updates, mode='drop')

[SNMS95]

Thank Jake!