Allow to specify batch dim during maybe_pad_batch (eg. for sharded datasets).
Allow to specify batch dim during maybe_pad_batch (eg. for sharded datasets).

PiperOrigin-RevId: 415495182

google-research/scenic