Parallelize different functions on independent inputs

[JamesYang007]

Hi! I have code that looks something like this:

import jax

def f_iter(args):
    x, y = args
    # parallel step
    x = some_func(x)
    y = different_func(y)
    # combine step
    x, y = mixing_func(x, y)
    return x, y

x, y = jax.lax.while_loop(
    some_cond_func,
    f_iter,
    some_init_val,
)

I would like to parallelize the computation of some_func and different_func on multiple cores because they operate on independent inputs. The non-trivial aspect is that at each iteration after the parallel step, we must mix the two values. Is the parallel step possible?