Variable sized batches

[pmrv]

We are trying to solve the following problem in grain. Our data for our JAX models requires postprocessing before it can be fed to the models and this postprocessing produces variable sized data. Because of this we batch and pad multiple input items together to achieve batches of the same size to avoid recompilation. The postprocessed data is much larger than the original input items.
Our problem is now that the Batched transformation for DataLoader in grain only allows to specify a fixed number of input items for a fixed batch size, whereas we would rather ingest a variable number of input items until we can reach a given output batch size with minimal padding. Batched doesn't seem to support this. Is this something that could be done with the DataSet API or do you see another way of mapping the problem onto the DataLoader?

[mhyatt000]

Not a maintainer, but I have found dataset API to be more flexible. More info about the problem might help determine best solution.

ie: how is the variable batch size determined? is it purely rng or is there some measurement you can use to batch intelligently

I also suffer from the same problem of too much padding, but in my case the data must all be shuffled together so I accept the additional padding. conversely, it sounds like you are willing to remove the padding at the cost of some randomness in the data order