Can functions with traced variable control flow be pmapped?

[xmax1]

Can I distribute a function with pmap, containing some jittable subfunctions, and a traced input dependent (not collective ie not dependent on the data from other devices AND not static) control flow at the end?

or do all pmapped functions have to only contain any code that is jittable?

Thanks!

docs https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html under heading 'pmap and jit'
control flow https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#control-flow
potentially related #5895

Example

def f(x):       # is this whole function pmappable? 
  x = x**2 + x                  # this is jittable
  if jnp.sum(x) < 3:           # this is not, do we have to retrieve the data from all the devices before running this? 
    return 3. * x ** 2
  else:
    return -4 * x

The docs indicate that no dynamic control flow can be used with pmap (by inference, they don't directly state this). Is there a way to distribute code with traced variables control flow?

[xmax1]

https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html

reading this closer, the answer really seems to be no.

[jakevdp]

Python control flow based on the values within a traced array cannot be used with JIT or pmap, but you can often re-express your code in such a way that it can. In your case, using jnp.where will let you express the same logic in a way that is compatible with JAX transforms such as jit and pmap:

def f(x):
  x = x**2 + x
  return jnp.where(jnp.sum(x) < 3, 3 * x ** 2, -4 * x)

There is more information available on this in 🔪 JAX - The Sharp Bits 🔪 : Control Flow

[xmax1]

Thanks for getting to this.

I see how where works in this instance. However, there is a particular reason why the control flow is needed.

def f(x):
x = x**2 + x
return jnp.where(jnp.sum(x) < 3, expensive_function(), cheap_function() )

I only want to call the expensive function for the whole batch on the worker gpu if at least one of the batch values is less than 3 (in this example).

As far as I undertand it, jnp.where will compute both the expensive and cheap function regardless of the condition?