[Hendrik-code]

What

Every GPU transfer augmentation (AugType.TA, 17 of them) can now blend its output back
towards the image it was handed, at an alpha drawn from a range set in the config:

"RandomScharrGPU": { "p": 0.25, "mix_prob": 0.5, "mix_alpha": [0.0, 0.4] },
"RandomSynthSegGPU": { "p": 1.0, "mix_prob": 0.7, "mix_alpha": 0.5 }

mix_prob is the existing per-sample gate for whether to blend; mix_alpha is new and
sets how far. It takes a scalar or a [lo, hi] range, mirroring
AnatomicalLabelOverlay.blend_strength.

alpha is the weight on the original: 0.0 is the transform's output, 1.0 is the
input untouched. Counter-intuitive in a config, and kept that way deliberately — it is
the direction the existing code used, which is what lets the default range (0.0, 1.0)
reproduce the previous hardcoded U(0, 1) draw exactly (draw * 1.0 + 0.0 is exact in
IEEE-754).

Why

Six of the seventeen already had half of this, from three open-coded copies of the same
four lines, with the weight hardcoded to U(0, 1) — so the only expressible blends were
"none" and "anywhere between the identity and full strength". The other eleven had no
blend at all. One shared implementation now lives in smauglab/transforms/gpu/mixing.py.

A configurable range also fixes a real hazard the hardcoded one had: inversion is
max - x, so alpha = 0.5 on RandomInverseGPU is max / 2 everywhere — a constant
image. U(0, 1) samples straight through that dead zone; [0.7, 1.0] does not.

Does it make sense for every TA?

Yes, for all seventeen. Every registered TA is intensity-only on a fixed voxel grid: none
resamples, deforms, flips or crops, none changes shape or channel count, none touches the
segmentation. The genuinely geometric transforms are tagged GEO, not TA, so there is
no augmentation in this group where a per-voxel convex combination would ghost or break
image/label alignment. RandomGaussianBlurGPU comes along because it shares
_RandomConvBaseGPU and already exposed mix_prob.

Three caveats are documented in the docstrings rather than designed around:

• RandomSynthSegGPU is the only one with no stats restore after the blend. Its
  synthesis is in [0, 1] while a training patch is z-scored, so there the blend really
  does mix two scales. Flagged in its Args, with a pointer to RandomPaletteGPU or
  ZscoreNormalizationGPU.
• RandomUnsharpMaskGPU / RandomBiasFieldGPU — the blend is close to a rescale of
  unsharp_amount / coefficients. Offered anyway: interpolating toward the identity is
  a different axis from scaling the effect, and uniformity is what makes a sweep writable.
• RandomPaletteGPU / PaletteSynthesisGPU blend in the [0, 1] space the synthesis
  is built in, before their closing foreground z-score, so both sides are in-units and the
  * flat_m_all masking still holds the background at exactly zero.

Where the blend sits

Before each transform's retain_stats restore, where the existing copies already had it.
That keeps the promise retain_stats makes — blending after the restore leaves the
output with the input's mean but a smaller standard deviation (s·sqrt((1+rho)/2), ~0.8s
at rho = 0.3). mix_prob is above zero in 29 of the 33 shipped configs, mostly on
transforms that ask for retain_stats too, so this is not hypothetical. Confirmed by
execution: with mix_alpha=0.75, retain_stats=True the output matches
restore(blend(...)) bit-exactly and differs from blend(restore(...)) by up to 1.36,
and the output's mean/std match the input's to 9e-08 / 0.

One behaviour change, measured

The two per-sample copies spent a draw per sample to decide something mix_prob had
already settled:

if torch.rand(1).item() < self.mix_prob:   # drawn even when mix_prob is 0.0

so RandomInverseGPU and RandomHistogramEqualizationGPU shifted the RNG stream for
every transform that ran after them, whether or not mixing was on. Both helpers now
return before drawing anything, as the vectorised copy always did.

Measured, by hashing every shipped config's seeded output before and after:

• 15 of the 28 runnable configs produce a different seeded sequence. Same
  distribution; a GPU augmentation run is not reproducible run-to-run under a fixed seed
  anyway, and no published result is recomputed from these files.
• Every transform in isolation is bit-identical except RandomInverseGPU with
  mix_prob > 0, whose blend also had to move ahead of its hoisted retain_stats.
• Isolated to contrast.py by reverting file-by-file, and attributed to
  RandomInverseGPU (6 draws → 4) with a torch.rand counter.

What it buys is the property this feature needs: a mix_alpha sweep is comparable
against its own mix_prob = 0 baseline at a fixed seed, which it was not before. Config
hashes are computed over the document, not over defaults, so the paper hashes the README
pins (c69b9872, 580beb13) and the ...-c-<hash> run folders are untouched.

Also here

Both found while wiring the above:

• RandomSynthSegGPU.apply_transform never cloned its input, writing input[:, c] = ...
  straight into the caller's batch["data"]. Every sibling already clones; it slipped
  through because test_transforms_gpu.py does not list that module. One-line fix.
• AnatomicalLabelOverlay.blend_strength now shares the range validator, so a value
  outside [0, 1] is rejected rather than merely documented as being in it.

Tests

unit_tests/test_mix_alpha.py, 19 cases. The ones that matter most iterate over
registry.entries(Backend.GPU, group=AugType.TA), so a TA added later cannot arrive
without the knob, or with it wired into the constructor but never read:

• the constructor accepts both, and the defaults are exactly 0.0 / (0.0, 1.0);
• mix_prob=1, mix_alpha=0.5 differs from mix_prob=0 — the guard against a dead knob;
• mix_alpha=1.0 is the identity, for 14 of the 17 (allclose: a retain_stats restore
  of the input's own stats carries its own 1e-8). The two palette classes are skipped
  with a reason — their unconditional z-score runs after the blend;
• a narrow range bounds every recovered alpha, two disjoint ranges stay disjoint, a scalar
  gives every sample the same alpha;
• mix_prob=0 is stream-neutral, and a degenerate range still spends its alpha draw, so
  an A/B over ranges stays comparable;
• the range survives the JSON path as a list, through validate_section and
  build_gpu_pipeline.

unit_tests/test_mix_prob.py passes unchanged — that is the check that mix_alpha's
default still reproduces the draw it replaced.

Full suite in seg11: 11 failed, 485 passed, against a pre-change baseline of
11 failed, 466 passed on the same interpreter — zero new failures. The 11 are
pre-existing: 9 test_trainers subtests (nnunetv2 2.4.2 here against the >=2.5 pin)
and 2 from the stale transform_params_gpu_random-normalization.json, which names five
transforms that are not registered. pre-commit run --all-files clean.

A second commit makes the two mix attributes visible to mypy. configure_mix assigned
them from outside the class, so every read resolved through nn.Module.__getattr__
(Tensor | Module) and the Mixing Protocol could not match — the exact problem it was
written to solve. They are declared on ImageOnlyTransform now, at the values that mean
"off"; only the transforms that take them as constructor arguments are configurable,
since the signature is the config schema. It also rejects a bool range, which the
numeric tower would otherwise have read as alpha 1.0. No behavioural change — every
per-transform digest is unchanged. All 8 CI checks green.

🤖 Generated with Claude Code