RandomMixUpV2 / RandomCutMixV2 cast class labels to the image dtype: bfloat16 corrupts labels above 256

Summary

RandomMixUpV2 and RandomCutMixV2 cast integer class labels to the image dtype. Under bfloat16 (and float16 above 2048) the labels themselves are rounded, so e.g. ImageNet-1k class ids come back wrong:

import torch, kornia.augmentation as K
img = torch.rand(4, 1, 8, 8).bfloat16()
out, lab = K.RandomMixUpV2(p=1., data_keys=["input", "class"])(img, torch.tensor([257, 999, 2049, 4097]))
print(lab[:, 0].tolist())   # [256.0, 1000.0, 2048.0, 4096.0]

float16 gives [257, 999, 2048, 4096]. Same for CutMix.

Expected

Labels should not share a lossy dtype with the image: keep the label columns in a dtype that represents them (e.g. float32/float64 floor, or return lambda separately).

Measured on main @ c643312a3 merged with #4634 (docs-only), torch 2.14.0, CPU, macOS arm64. Found while reviewing #4634 (batch 6d of #4407).

Posted on behalf of @ducha-aiki by Claude (Fable 5.1).