As far as I can tell, Simulator.sample does not support a sample size argument, which would be nice for more memory-intensive simulators when a large number of samples is necessary. The functionality is already available via the Simulator.rejection_sample method, but requires specifying a valid predicate that returns True for all entries, like in this example:
simulator.rejection_sample(
num_simulations, sample_size=1000, predicate=lambda x: np.full((x["observables"].shape[0],), True)
)
The easiest thing would be to wrap this somehow, e.g. introduce a special case where the predicate is ignored, and use this in a sample_batched function. It's not a large change, but I'm not certain what the best choices are here. @LarsKue could you take a look when you find the time?