[Hendrik-code]

What does this change?

Adds the five normalization schemes the GPU side was missing, and makes the choice
between them an augmentation.

New in smauglab/transforms/gpu/normalization.py:

class	scheme	nnU-Net equivalent
MinMaxNormalizationGPU	rescale the range onto out_range, optionally after a percentile clip	RescaleTo01Normalization at defaults
PercentileNormalizationGPU	clip to a percentile window or an explicit clip_range, then standardize	CTNormalization given clip_range, mean, std
RobustZscoreNormalizationGPU	median, scaled by the IQR or the MAD	—
MaxScaleNormalizationGPU	divide by max|x|, or by a high percentile of it	—
MeanScaleNormalizationGPU	divide by the (masked) mean	—

All five share a base class that owns the clone, the channel guard, the statistics mask
(stats_from: all voxels, the labelled anatomy, or everything preprocessing did not
zero), both degeneracy guards and the in_seg/out_seg region selection, so each scheme
is an affine map plus an optional clip.

The choice reuses RandomChooseXTransformsGPU rather than adding a combinator. A new
AugType.NORM group collects the normalizations, and the builder emits them as one
num_transforms=1 bucket at the end of the pipeline in every pipeline mode. Because
that bucket draws one child uniformly and then spends that child's own p, each scheme's
p becomes its weight (P(scheme) = p / n) and the remainder is "no normalization at
all" — so there is no no-op transform to register. pipeline.random_choose.norm_probability
gates the stage.

Why?

Every scheme other than z-score lived only in nnU-Net's preprocessing, which runs once per
dataset. Leaving the intensity standardization fixed is one more thing a model can come to
depend on, and it is the step that differs most between sites and pipelines, so it is worth
randomizing for the same reason the contrast and gamma augmentations here are randomized.

How was it tested?

• pytest passes locally
• pre-commit run --all-files passes
• Added or updated tests covering the change
• Ran a training / augmentation job end to end

unit_tests/test_normalization.py (33 tests) covers each scheme's maths, exact parity
against the installed nnU-Net's RescaleTo01Normalization, CTNormalization and
ZScoreNormalization, both guards, the masked statistics, the builder's staging in all
three pipeline modes, and the weight semantics as a 400-draw statistical test. The new
classes also join the registry-driven suites (test_transforms_gpu.py,
test_container_fast_path.py, test_no_inplace_mutation.py, test_channel_guard.py).

Full suite, nnunetv2 2.4.2 env: 9 failed, 501 passed — the same 9 test_trainers.py
subtests that fail on origin/main in that env (order_resampling_data / border_val_seg,
the documented >=2.5 pin). Verified against a clean origin/main worktree: 9 failed, 468
passed. Also ran the normalization, builder, channel-guard, transform and registry tests
under nnunetv2 2.6.2; the only failures there are the two pre-existing
batchgeneratorsv2 SpatialTransform ones, confirmed identical on clean main.

End to end on one GPU: 50 batches of (2, 1, 64, 64, 64) through a pipeline whose
normalization is a six-way choice, shapes, finiteness and label integrity held, and
record_applications shows all six schemes firing in proportion to their p.

Anything reviewers should look at closely?

• min_scale is measured against the standard deviation, not the intensity range.
  The range was the first thing tried and is wrong for the robust schemes: one 1e4 voxel
  puts the range four orders of magnitude above the IQR, so a range-relative floor
  declined to run on exactly the volume RobustZscoreNormalizationGPU exists for.
• One draw per batch, not per sample, at the shipped defaults. pipeline.same_on_batch
  defaults to True and the container forces it onto every stage it holds, this bucket
  included — the same is already true of the ta and ge buckets. pipeline.same_on_batch: false gives each sample its own scheme. Pinned in TestOneDrawPerBatchOrPerSample.
• This randomizes the final input scale on purpose. The patch arrives already
  normalized by preprocessing and inference re-applies only that, so a config turning this
  on is deliberately training scale invariance rather than matching inference. There is no
  compensating re-normalization at the end; that was a conscious choice, and it is written
  up in the README section.
• PercentileNormalizationGPU takes mean/std/clip_range as parameters because
  nnU-Net reads them from the dataset fingerprint and nothing hands a transform that at
  training time. A config that knows its dataset states them; left as None they are
  measured per sample.

Checklist

• Augmentation behaviour is unchanged, or the change is intentional and described above

  One intentional change: ZscoreNormalizationGPU and the CPU ZscoreNormalization move
  from AugType.GE to AugType.NORM, which is what puts z-score into the same choice as
  the new schemes. In random_order mode that takes it out of the shuffled GE bucket and
  makes it a trailing stage. Every shipped config sets its p to 0.0, so nothing fires
  differently — but the GE bucket no longer spends a uniform draw on it, so random_order
  runs will not reproduce earlier volumes bit for bit under the same seed. A config naming
  a single normalization is otherwise left exactly where it was, including under
  pipeline.order: config.

• New transforms are reachable from a config JSON

  smauglab/configs/all_augmentations.json and the README matrix are regenerated with
  smauglab matrix --write.

🤖 Generated with Claude Code