Add warning for varying simulator output sizes

[LarsKue]

Varying simulator output sizes are a common occurrence when the number of samples varies between calls to simulator.sample():

def context(batch_size):
    n = np.random.randint(10, 101)
    return dict(n=n)

def prior():
    mu = np.random.normal()
    sigma = np.random.gamma(shape=2)
    return dict(mu=mu, sigma=sigma)

def likelihood(n, mu, sigma):
    y = np.random.normal(mu, sigma, size=n)
    return dict(y=y)

simulator = bf.make_simulator([prior, likelihood], meta_fn=context)

However, these can trigger excessive compile times in JAX, where each value for n triggers a recompilation. For a wide range of n, this can mean that the compilation dominates the training time.

The current best-practice fix for users is to use padded tensors:

def likelihood(n, mu, sigma):
    y = np.random.normal(mu, sigma, size=100)  # uses fixed maximum size
    y[n:] = 0  # set unused entries to zero, or some other placeholder value
    return dict(y=y)

When we detect that compile times dominate, we should output a warning to the user, with a suggested fix. We could also improve support for padded simulator output in general. Further, we could look into if there are better ways to mask out unused values rather than just setting them to placeholder values like above.

[paul-buerkner]

It sounds as if padding could be a great adapter feature, something like

adapter.pad(variable_dict, len = 100, axis = 1, value = 0)

I wouldn't want to burden the simulator with padding, since the simulator describes the probabilistic program that we would ideally keep free of any technical stuff related to deep learning. Sure, we could code padding within the simulator as a user, but I would prefer having an adapter functionality that is easier to get right for the user and doesn't mess with the probabilistic program.

[LarsKue]

Closing this as users can just switch backends. Reopening with padding as feature request.