Probabilistically merge two generators in Jax

[aeftimia]

Say I have two generators whose elements consist of a tuple of (jax array, metadata). To a first approximation, I'd like to construct "merge" operation that creates a new generator as follows:

1. Create a queue with the first element of generator1 and generator2.
2. Compute p = <some pure function>(queue[0], queue[1])
3. i~bernoulli(p)
4. yield queue[i], p
5. update queue with next element of the appropriate generator

I can't think of a good way to do with without hitting side effects. I looked into scan as well, but it's still not clear it's relevant.

Thanks,
Alex

[jakevdp]

When you say "generator" what do you mean? Is it literally a Python generator object?

[aeftimia]

@jakevdp, correct, that's my ideal situation. I'm ultimately looking to differentiate sum(log(p)) for each merge over many nested merges with respect to some parameters passed through metadata.

[jakevdp]

In that case, your operation is inherently side-effecting (generators are side-effecting, after all), so there's no way to handle such a thing in transformable JAX code. The best approach would probably be to write a Python function that merges the two generators and outputs a new generator. Then you could pass the output of that generator to JAX.

[aeftimia]

Ah well. So be it!