Doing a reduce on arange without materializing the vector

[mgostIH]

@partial(jax.jit, static_argnums=[0])
def do_funny_thing(N):
    v = jnp.arange(N)
    return jax.lax.reduce(v, 0, lambda x, y: x ^ y, [0])

do_funny_thing(1_000_000_000_000)

The following code goes out of memory, it would be lovely if it wouldn't. Is there some known method to avoid materializing an arange, and more generally any kind of vector that is defined in terms of it, like those generated from jax.random? I'm not sure whether it's the reduce actually materializing it in memory, if that's the case what should I do?

[YouJiacheng]

You can try lax.while_loop. Note this might be slow, you can loop over chunks instead of elements.
alternatively, you can convert [i for i in range(1000)] into [i * 1000 + j for i in range(1000) for j in range(1000)], i.e. a nested scan on two arange.
see #10401 (reply in thread)

[mgostIH]

Indeed, I thought about the while loop but I'd imagine this would be extremely slow due to its sequential nature, while a lot of the tasks I want to do would be based on reduce.

[YouJiacheng]

So you should not process only one element per iteration in while loop!

jax.lax.fori_loop(0, 1_000_0000, lambda i, val: jax.lax.reduce(jnp.arange(i * 1_000_000, (i + 1) * 1_000_000), val, lambda x, y: x ^ y, [0]), 0)

[mgostIH]

I will consider this as a possible alternative, but I thought that not having to load anything on device and not needing to think about the loop splitting would be a better design for trivially parallelizable tasks.