Scan an array and conditionally append

[mmuckley]

Hello. I have an application where I would like to scan over an array and, depending on the current value and state, conditionally append to an output list. A single value in the input array could result in zero to many append calls. I have an idea of the maximum size of the output array, but in practice it will usually be much smaller, and I can't predict in advance how small.

For my current implementation I pre-allocate the output array and insert values using jax.ops.index_update, but this copies the entire output array at every value, which is extremely slow.

I was wondering if there could be a way to do this using a pytree, but I don't really understand them. In pure Python, I would be able to call a.append(b), but since JAX uses pure functions, I can't do this as it mutates a.

Altering lax.scan would help if I could control the ys.append(y) line, but as I said earlier, a single value of x could result in zero or many appends.

Based on the recent update in sharp bits, I thought that it might be possible to get jax.ops.index_update by removing all references to input variables, but I didn't have any success with this. It could be that the index_update command is called deep in the stack and there is a dangling reference.

So I'm a bit at a loss for any way to speed up the code while using JAX and was curious if any others had any ideas. I don't think I can provide the exact code yet, but this might get at the idea:

import numpy as np

rng = np.random.default_rng(0)
data = rng.integers(0, 5, 50)
output = []
state = 0

for val in data:
    state += val
    if state >= 5:
        while state > 0:
            output.append(state // 2)
            state = state // 2
        state = 0

[davisyoshida]

I'm not sure if this is what you've already done, since you said you tried using index_update, but here's a way to convert the toy example to use scan. (Also I'm not a JAX expert, so there may totally be a better way to do this). It also uses a copy (the write_array.at... bit), but I think if you jit it, the JIT will know that the intermediate arrays aren't actually used so it should be able to do the write in place.

key = jax.random.PRNGKey(0)
data = jax.random.randint(key, (50,), 0, 5)

def while_cond(carry):
    _, _, value = carry
    return value > 0

def while_body(carry):
    write_array, write_index, value = carry

    value = value // 2
    write_array = write_array.at[write_index].set(value)
    return write_array, write_index + 1, value

def f(carry, val):
    write_array, write_index = carry

    write_array, write_index, _ = jax.lax.while_loop(
        cond_fun=while_cond,
        body_fun=while_body,
        init_val=(write_array, write_index, val)
    )
    return (write_array, write_index), None

max_len = 150
init_carry = (-1 * jnp.ones(max_len, dtype=jnp.int32)), 0
(out_array, n_valid), _ = jax.lax.scan(f, init_carry, data)

[mmuckley]

Thanks for the suggestion @davisyoshida. Unfortunately this does not work for me, and I observe the same code speed. The documentation currently states that index_update is aliased to the at commands.

Currently, my insertion operation is fairly deep in a function stack and inside a loop, so it is possible there is a dangling reference that is preventing jit from doing this in-place. I have tried to remove all references, but it doesn't speed up the code. It would be nice if there were a way to force this to be in-place.

[C-J-Cundy]

Quick couple of comments (I'm also not a jax expert!):

1. You might want to look at jax.lax.dynamic_slice_in_dim. I've had some success using this where the index_update functions fail. (That was quite a few jax updates ago though)
2. I'd recommend thinking about if it's possible to change the overall code so that the output you need from this function isn't a list/array of variable length. XLA/JAX really wants to know the shapes of things if you're going to use jit, so it's going to be a struggle to jit this function. For instance, maybe the output could be an n_data x max_per_data array with another n_data x max_per_data mask array to show where the outputs are from each data point?

[mmuckley]

1. Yes I found this function as well. Unfortunately, I didn't have any luck with it. I've also tried making a smaller buffer that I periodically write into the big array with slice updates, but that also didn't work.
2. Are you saying that max_per_data is the max for each element of the input array? This might work, but is extremely memory greedy. On average the output will be very sparse, but in the (unlikely) worst-case scenario it can be very big. It might be better to just support the worst case scenario with slow code than have fast code that fails randomly.

My other thought is that with lax.scan I thought the machinery might already be in XLA to support this, but it hasn't been exposed.

[C-J-Cundy]

Yes, I was assuming that there is some maximum output which each input item can result in, max_per_data. I'd recommend trying to vmap over the inputs, assuming the worst case.

I guess in your application the outputs must be quite large or you must have a huge amount of data?
If you were writing just a single float per output you could have max_per_data be 1000 and with 100,000 inputs that would only be 400MB.

[jakevdp]

The fundamental difficulty here is that the output of your code is an array whose shape depends on the values of the inputs; this is incompatible with JAX's tracing and compilation model, which requires the shapes of input and output arrays to be known statically.

Regarding index_update: it is true that in normal op-by-op mode it leads to a copy of the array each time it is called, but when JIT-compiled XLA can optimize away such copies. If you have a simple example of an index_update approach where this optimization appears not to be happening, we might be able to help diagnose that issue.

[mmuckley]

Okay I will try to construct such an example. However it's difficult to diagnose at a high level whether the in-place update is happening. Having an option that would force the in-place update or cause the program to fail to compile would help greatly.

[schwarzTass]

Has anyone made progress on this question? Even trying the jax.lax.dynamic_slice_in_dim command as @C-J-Cundy suggested is very slow, suggesting that this is just copying as well.

I'm in the same situation of @mmuckley where the output is very sparse, only every say 1000th step on average is relevant information; even though I know how much relevant information I will have in the end, just outputting at every step and then filtering this very sparse array consumes a lot of memory.

A minimal example could look like this:

import jax
import jax.lax as lax
import jax.numpy as jnp


input_array = jnp.arange(10000)
every_1000th_result = jnp.zeros((input_array.shape[0] // 1000,3))

@jax.jit
def scan_body(carry, input_val):
    every_1000th_result, next_idx_to_write,prev_result,iteration_number = carry
    
    # update the result
    current_result = prev_result + input_val

    # only if iteration_number is a multiple of 1000, write to every_1000th_result at index next_idx_to_write
    every_1000th_result = lax.cond(
        iteration_number % 1000 == 0,
        lambda x: x.at[next_idx_to_write,:].set(jnp.array([current_result,current_result+1,current_result+2])),
        lambda x: x,
        every_1000th_result
    )
    # increment next_idx_to_write
    next_idx_to_write = lax.cond(
        iteration_number % 1000 == 0,
        lambda x: x + 1,
        lambda x: x,
        next_idx_to_write
    )
    # increment iteration_number
    iteration_number = iteration_number + 1

    # return the updated carry and the current result
    return (every_1000th_result, next_idx_to_write, current_result, iteration_number), None

# Initialize the carry
initial_carry = (every_1000th_result, 0, 0, 0)
# Use lax.scan to iterate over the input array
final_carry, _ = lax.scan(scan_body, initial_carry, input_array)
# Extract the result from the final carry
print("Got every_1000th result:", final_carry[0])

Here, I want to make sure that x.at... does not create a copy of the entire array every_1000th_result. How can this be done?

[NB: I'm returning None from my scan body(as proposed by @davisyoshida), and do not use the second result of the jax scan output. Does jjit compile this away and understand that it should not save a copy of 10000 nones?]

[patrick-kidger]

FWIW I find your script to run efficiently on my machine (In about 10^-5 seconds), but if you need it then a trick to avoid copying big arrays in conditionals is to rewrite:

lax.cond(predicate, lambda x: x.at[idx].set(v), lambda x: x, big_array)

into:

v = jnp.where(predicate, v, big_array[idx])
big_array.at[idx].set(v)

I believe XLA will actually already do the above rewrite under some scenarios (I contributed the compiler pass that does this 😁 ), although I don't know if it's perfect -- perhaps it misses some scenarios in which it is beneficial for it to do so.

[schwarzTass]

Thank you, that helps a lot and does the trick for my case! May I just ask for a better understanding why yours avoids the copies of big arrays? Because in principle, my understanding is that both versions may be compiled down to a select() call right (at least if they were to be vmapped)? Does you solution avoid the big copies because the .at happens outside of any condition?

PS: You are right, the minimal example above also works fast on my machine, but if things get larger, I see a big difference :).

[patrick-kidger]

If you consider the first snippet in my response, you can see that evaluating this naively under vmap (so that cond becomes select) then it would involve copying big_array. You need to materialise both branches to pass in to select, both the version with the update and the version without!

The second snippet I have makes sure to avoid this problem: as you say, the .at happens outside of any condition.