Automatic Batching using `vmap`

[LarsKue]

Classes such as bayesflow.simulation.ContextGenerator or bayesflow.simulation.Prior allow the user to pass in a batched or an unbatched generation function, where typically, passing both is disallowed.

For quality of life (and possibly compatibility with #123) it would make sense to change this to a simple boolean flag, and automatically parallelize unbatched functions for the user with a variation of either JAX's or PyTorch's vmap function.

Example pseudo-code:

Replace this

class Prior:
    def __init__(self, batch_prior_fun=None, prior_fun=None):
        if (batch_prior_fun is None) == (prior_fun is None):
            raise ValueError

        self.prior = prior_fun
        self.batched_prior = batch_prior_fun

with

class Prior:
    def __init__(self, prior_fun, is_batched=True):
        self.prior = prior_fun if is_batched else vmap(prior_fun)

Since this is a breaking change, we would need a soft introduction of this feature including backward compatibility and deprecation messages.

Possible concerns:

• vmap is backend-dependent, and not always compatible with all operations that could happen in prior_fun (e.g. data-dependent control flow), so allowing the user to disable this would be beneficial
• sequential generation of batches is inherently incompatible with vmap, since vmap makes the stronger assumption that each batch item is independently processed

[LarsKue]

Implemented with 6565a51a9011554d3374543bdcd454f86878c5d8