Auto policies cannot be deep-copied after a forward: OperationBase stores an unused non-leaf Bernoulli sampler (learned probability is inert)

Summary

#4469 tracks that the auto policies cannot be pickled. copy.deepcopy fails too, for a different reason, and only after the policy has been used:

import copy, torch; from kornia.augmentation.auto import RandAugment, AutoAugment, TrivialAugment
x = torch.rand(4, 3, 32, 32)
r = RandAugment(2, 15); copy.deepcopy(r)     # OK
r(x); copy.deepcopy(r)
# RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment

AutoAugment() / TrivialAugment() fail the same way after .train() / .eval(), and so does K.AugmentationSequential(RandAugment(2, 15)) after a forward. A forward under torch.no_grad() does not trigger it.

Cause

OperationBase._update_probability_gen stores Bernoulli(self.probability) / RelaxedBernoulli(...) on op._p_gen / op._p_batch_gen; self.probability is a non-leaf clamp of an nn.Parameter. Those samplers are never read: grep -rn "_p_gen\|_p_batch_gen" kornia/ finds only the four assignments, and kornia/augmentation/base.py gates on the Python float op.p (already noted in #4428). Consequences measured:

from kornia.augmentation.auto.operations import ops
for raw in (1e-7, 0.5, 0.9999999):
    o = ops.Invert(initial_probability=0.5); o._probability.data.fill_(raw)
    print(raw, float(o.forward_parameters(torch.Size([4000, 1, 4, 4]))["batch_prob"].mean()))   # 0.49, 0.50, 0.50
o = ops.Brightness(initial_magnitude=0.3, initial_probability=0.5); o(torch.rand(8, 3, 4, 4)).sum().backward()
print(o._probability.grad, o._magnitude.grad)   # None  tensor([49.])

So the "learnable probability" is inert (no effect on sampling, no gradient), and its only observable effect is breaking deepcopy.

Expected

Either wire the sampler into the gate, or stop building it; in both cases do not store a non-leaf tensor on the module.

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