Ensure mixup in the imagenet data pipeline is getting a new random seed per batch.
Ensure mixup in the imagenet data pipeline is getting a new random seed per batch.

Before, a `functools.partial` set mixup to always use the same `rng`. Now we correctly fold in a unique batch index to produce a `per_batch_mixup_rng`.

PiperOrigin-RevId: 445138108

google/init2winit