RandomMixUpV2 / RandomCutMixV2 apply `p` twice: a batch-wide gate and a per-row Bernoulli, so rows are mixed at rate p²

Summary

RandomMixUpV2(p=...) and RandomCutMixV2(p=...) map p to the base batch-wide gate (p=1.0, p_batch=p), and then their generators draw a second, per-row Bernoulli(p) that zeroes mixup_lambdas / the cut size. A row is therefore mixed with probability p², not p, and inside a "selected" batch individual rows are silently left unmixed.

Reproduction

import torch, kornia.augmentation as K
x = torch.rand(4, 1, 8, 8); sel = rows = mixed = 0
for s in range(400):
    torch.manual_seed(s); a = K.RandomMixUpV2(p=.5); a(x)
    bp = a._params["batch_prob"]; assert bp.unique().numel() == 1   # batch-wide gate
    if bp[0] > 0:
        sel += 1; rows += 4; mixed += int((a._params["mixup_lambdas"] > 0).sum())
print(sel, mixed, rows)   # 195 selected batches; 405 / 780 rows actually mixed

Same loop with RandomCutMixV2(p=.5, use_correct_lambda=True): 295 / 780 rows changed inside selected batches. Per-row rate ≈ 0.5 · 0.5 = 0.25 rather than the requested 0.5. At p=1 both draws are always true, which is why no test sees it.

Source: random_generator/_2d/mixup.py (mixup_lambdas = mixup_lambdas * batch_probs) and the matching cutmix.py generator, on top of p_batch=p in the class constructors.

Expected

p applied once. Either the batch gate or the per-row draw, not both.

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