Mix augmentations raise for an unsupported data key only when a sample was selected; at p=0 the key passes through silently

Summary

A mix augmentation asked for a data key it does not implement raises NotImplementedError only if at least one sample was selected. When the gate selects nothing (p=0, or an unlucky draw at any 0<p<1) the annotation is passed through silently. The same pipeline therefore raises on some batches and not on others.

#4493 (closed) fixed the silent pass-through inside AugmentationSequential; this is the direct-call twin.

Reproduction

import torch, kornia.augmentation as K
img = torch.rand(4, 1, 8, 8)
for p in (1.0, 0.0):
    for make, ann, key in (
        (lambda: K.RandomJigsaw(p=p, grid=(2, 2)), torch.arange(4), "class"),
        (lambda: K.RandomMixUpV2(p=p), torch.rand(4, 1, 4), "bbox_xyxy"),
        (lambda: K.PatchMix(p=p, patch_size=2), torch.rand(4, 2, 2), "keypoints"),
    ):
        try: make()(img, ann, data_keys=["input", key]); print(p, key, "silent")
        except NotImplementedError: print(p, key, "NotImplementedError")
# p=1.0: all three raise.   p=0.0: all three silent.

Full matrix at p=0: boxes/keypoints pass silently on all five classes; class/label pass silently on PatchMix, Jigsaw and Mosaic; only mask raises regardless. Cause: transform_boxes / transform_keypoint / transform_class in _2d/mix/base.py reach the raising handler only under sum(to_apply) != 0.

Related, same gate: RandomJigsaw(grid=(2, 3), p=0.)(torch.rand(2, 1, 9, 13)) returns normally although 9x13 is not divisible by the grid; at p=1 it raises RuntimeError. At the default p=0.5 both failures are draw-dependent.

Expected

Support for a key (and grid divisibility) is a property of the class and the input, not of the draw: validate before the gate.

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).