add a minimal implementation of batching without sampling

[KathyHaem]

• keeps the original batching implementation around
• changes the default settings to apply standard batching

Limitations:

• I ran into a problem at the end while re-loading the sharded FLAX model. Will work fine if the model is small enough to be saved as a single file. The predictions are saved in FLAX format regardless
• This still seems to take a large amount of CPU memory, in particular there was a spike at the end of the runtime (I assume once inference is done and the model is being saved).

5

[KathyHaem]

@bminixhofer Did you see this/get a chance to take a look?

[bminixhofer]

Sorry I was swamped the last week but will test this as soon as possible. Thanks for the PR!