SeqIO evaluation hangs on final batch when dataset is padded

Hi all,

During SeqIO inference evaluation, the data pipeline is hanging for a long period of time on the final batch of data **only when the dataset is padded**. Sometimes for hours depending on the dataset size. This relates to these [lines of code](https://github.com/google-research/t5x/blob/7a10c1ff5d2ef5714c3b4d32192e5ee6b843d56f/t5x/utils.py#L1126-L1146).

This is not an issue when the dataset is not padded, i.e. when the dataset length is divisible by batch size.

I am guessing it is the `pad_ds` on this [line](https://github.com/google-research/t5x/blob/7a10c1ff5d2ef5714c3b4d32192e5ee6b843d56f/t5x/utils.py#L1144-L1145) that is very slow to process, but I am not sure why:

```python
pad_ds = ds.take(1).map(lambda i, x: (np.int64(-1), x)).repeat(
dataset_pad_amt)
```

Any idea how we could speed up the data pipeline when padding?

Start in t5x/utils.py at lines 1126-1146 and reproduce SeqIO inference evaluation with a padded dataset, comparing it with a length divisible by the batch size. Investigate the pad_ds construction and confirm the final batch completes without the prolonged hang while preserving evaluation behavior.

Written by the indexing model from the issue text.