How can I implement divide and conquer functions, and will the independent parts be run in parallel (on GPU)

[AdrienCorenflos]

Hi,

I have a use case where I want to implement a recursive function which would look similar to the below:

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


def recursive_sum(a):
    print(a)
    stop_cond = a.shape[0] <= 1
    def _stop_fun(op):
        return op.sum()
    
    def _continue_fun(op):
        k = op.shape[0]
        return recursive_sum(op[:k//2]) + recursive_sum(op[k//2:])
    
    return lax.cond(stop_cond, _stop_fun, _continue_fun, a)

This however does not finish the way I would have expected. Is there a way for me to adapt the code to make it work? And would the splitted operations then run in parallel? I think the most similar thing that exists in the code base is the associative_scan which seems to not have a lazy condition. Should I make my cond not lazy?

Note that my real use case actually uses the switch functionality, not cond
So it would actually look more like the below

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


def recursive_sum(a):
    print(a)
    stop_cond = a.shape[0] <= 1
    def _stop_fun(op):
        return op.sum()
    
    def _continue_fun(op):
        k = op.shape[0]
        return recursive_sum(op[:k//2]) + recursive_sum(op[k//2:])
    
    return lax.switch(1 - jnp.int32(stop_cond), [_stop_fun, _continue_fun], a)

[jakevdp]

The issue here is that lax.cond evaluates both branches at trace time, and so as written the tracing will never terminate. You might be able to do something like this using lax.while_loop to put the stopping condition in the compiled code.

[davisyoshida]

Wouldn't while_loop not like the variably shaped slices required?