Improve slicing of padded input batch
Improve slicing of padded input batch

If reshard_inputs_fn (partially) replicates the sliced batch dim, a halo exchange will be triggered, and it could be a lot of collective-permute ops when the mesh is large. So we reshape the batch into 2d and perform the slicing in two steps.

PiperOrigin-RevId: 455524384

google/paxml,google/paxml