[JAX SC] required_buffer_size_per_sc is being undercounted when minibatching is enabled. The initial calculation does not account for the alignment padding added at the end of each minibatch.

[copybara-service]

[JAX SC] required_buffer_size_per_sc is being undercounted when minibatching is enabled. The initial calculation does not account for the alignment padding added at the end of each minibatch.

This potentially leads to the following sequence of events:

1. A smaller-than-needed buffer size was calculated and reported. FDO used this incorrect value to reconfigure the buffer for the next run.
2. The buffer would inevitably overflow and drop IDs.
3. Could cause a persistent FDO loop.

Now, we can recompute the required buffer size after the minibatches have been merged and their final memory layout is known.