feat: restore a jittered intensity profile around the whole GPU pipeline

[Hendrik-code]

What does this change?

pipeline.retain_stats brackets the augmentation stack with one snapshot and one restore
of the image's intensity profile, in place of the per-transform retain_stats flag:

[geometry] [snapshot] [augmentations] [restore] [crop] [normalization]

The snapshot records the per-sample, per-channel mean, standard deviation, 1st and 99th
percentile after the geometry. The restore runs before the trailing crop and
normalization.

Two of the four, drawn per batch. An affine map y = ax + b has two degrees of
freedom, so four targets over-determine it. Each pass therefore draws two of the four
and solves for the map that puts exactly those two on target, leaving the other two where
the augmentations left them. Writing each statistic as t = a * u + b * c, with c = 1
for the three that move with an offset and c = 0 for the offset-invariant standard
deviation, makes all six pairs one 2x2 solve:

drawn pair	a	b
mean, std	s'/s_x	m - a*m_x
mean, p1	(m - a')/(m_x - a_x)	m - a*m_x
mean, p99	(b' - m)/(b_x - m_x)	m - a*m_x
std, p1	s'/s_x	a' - a*a_x
std, p99	s'/s_x	b' - a*b_x
p1, p99	(b' - a')/(b_x - a_x)	a' - a*a_x

The drawn pair is itself an augmentation, and the pool spans the two useful extremes:
(mean, std) is what per-transform retain_stats did, and (p1, p99) is the robust dynamic
range, which the outliers RandomScharrGPU or RandomRandConvGPU manufacture cannot
move -- those outliers are exactly what make the standard deviation a bad target.
stats is the pool, so ["mean", "std"] reproduces the old behaviour at pipeline level.

Jitter. std_scale multiplies the standard-deviation target and percentile_scale
multiplies each percentile's distance from the mean, independently, so the span can widen,
narrow or go asymmetric. The mean target is never jittered: it is the anchor every pair
shares, and moving it too would put two unrelated augmentations behind one knob.

"pipeline": {
  "retain_stats": {
    "stats": ["mean", "std", "p1", "p99"],
    "percentiles": [1.0, 99.0],
    "std_scale": [0.9, 1.1],
    "percentile_scale": [0.9, 1.1],
    "stats_from": "all",
    "p": 1.0
  }
}

An absent block is off and emits nothing; a present but empty one is on with the defaults,
which restores the drawn pair exactly.

Why?

retain_stats restores once per transform, so a chain of them pays for N restores to end
up where the last one puts it; it knows only two statistics, so an edge filter's outliers
drag the standard deviation and the restore with it; and it restores the input's profile
exactly, which teaches the network that its input scale is reliable. That is the
opposite of what the rest of this pipeline is for. #135 says of its own randomized
normalization that "there is no compensating re-normalization at the end"; this is that
compensation, and it is placed before the normalization stage rather than after it.

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_pipeline_stats.py, 38 tests / 42 subtests: each of the six pairs lands
its two statistics (max |d| 0 to 6e-8 on a [2, 1, 24, 24, 24] batch) while the other two
demonstrably move; parity with _restore_stats when the pool is fixed to (mean, std);
the jitter stays inside its configured range over 25 seeds and is wider than 0.1; the draw
count is exactly five uniforms per channel and does not change with the ranges or the size
of the pool; a constant patch and an unlabelled mask are left alone; stats_from="seg"
really measures the anatomy only; the emitted order in all three pipeline modes and under
pipeline.order: config; that the only transform the bracket moves is the crop; the mask,
in-place and both-kornia-routes contracts; and the config surface, including that the two
paper configs do not carry the block.

Full suite, seg11 (nnunetv2 2.4.2): 9 failed, 621 passed, 18 skipped. Measured on
this branch's merge base in a worktree: 9 failed, 583 passed -- the same nine
test_trainers.py subtests, which are the documented >=2.5 pin, so this adds 38 tests
and no failures. ruff check / ruff format --check clean, mypy clean on the changed
files locally (2.1.0 here; CI's 2.4.0 is the real gate), smauglab matrix --check exit 0.

End to end on one A40, through smauglab/configs/transform_params_gpu_stats-restore.json
under torch.autocast("cuda", float16): 50 batches of (2, 1, 64, 64, 64) held shape,
device, finiteness and label integrity; record_applications shows the bracket around the
augmentations that fired; on a 4-channel batch with the pool fixed to (mean, std) the
restore lands within 3e-7 per channel while p1/p99 move; and a config that names a crop
emits ... StatsRestoreGPU, RandomCropTransformGPU, ZscoreNormalizationGPU and runs finite.

Anything reviewers should look at closely?

• Crop moves to the tail when the bracket is on. It is AugType.GEO but it is not an
  intensity augmentation: it occludes a random box with the volume's own minimum without
  changing the shape, so a profile measured before it would be restored against the
  occlusion. It is active (p: 0.25) in transform_params_gpu.json,
  transform_params_gpu_default01-23_CropTransform.json and
  transform_params_gpu_palette-em.json; none of them enables the bracket, and with the
  block absent nothing moves. test_the_only_transform_the_bracket_moves_is_the_crop
  pins that the rest of the order is untouched.
• A lone normalization also moves to the tail when the bracket is on. Without the
  bracket a single normalization keeps its pipeline position verbatim (feat: add the main GPU normalization schemes and a random choice between them #135's rule); with
  it, the restore has to come first, so it is staged. In random_order mode the snapshot
  deliberately lands before the hoisted RandomLowResTransformGPU, which is a general
  enhancement that merely runs early -- blurring moves the standard deviation and belongs
  inside the bracket.
• The two leaves share a mutable slot (gpu/stats.py), which the restore consumes and
  clears. The restore raises rather than guessing if it runs without a snapshot or if the
  batch size changed between the two. Both are pinned.
• Draw discipline. Five uniforms per channel, always, whichever pair comes out and
  whatever the ranges -- so a jitter sweep is comparable against its own baseline at a
  fixed seed, which is invariant 2 of gpu/regions.py. The pair is drawn by scaling one
  uniform into the table, the fixed-width trick UNIT_DRAW_WIDTH uses, rather than by
  torch.randint, so adding a statistic to stats does not reshape the stream.
• Population standard deviation (unbiased=False), matching ZscoreNormalizationGPU,
  nnU-Net and gpu/normalization.py -- not the sample one _channel_stats uses. On a 24^3
  patch that is a factor of 1.00004, which is why the parity test against _restore_stats
  has a tolerance.
• retain_stats is untouched, as requested. The new config sets it to false on all
  15 transforms that accept it, because the two do not compose. No existing config is
  modified: all 33 are byte-identical to this branch's merge base.
• This PR merges feat: add the main GPU normalization schemes and a random choice between them #135 (hm/gpu-normalization-schemes) into Finish in_seg / out_seg: where an augmentation applies, at what granularity, and how sharply it stops #138's branch, to reuse
  AugType.NORM, _normalization_stage, _masked_quantiles and _resolve_mask rather
  than write second copies. The only merge conflict was
  smauglab/configs/all_augmentations.json, resolved by regenerating it with
  smauglab matrix --write. So the diff shown here is a 4-deep stack: Give every GPU transfer augmentation a configurable mix range #137 -> Finish in_seg / out_seg: where an augmentation applies, at what granularity, and how sharply it stops #138 ->
  merged feat: add the main GPU normalization schemes and a random choice between them #135 -> this.

Checklist

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

  Unchanged with the block absent, which is every existing config.

• New transforms are reachable from a config JSON

  Not transforms: StatsSnapshotGPU and StatsRestoreGPU are pipeline structure, emitted
  by build_gpu_pipeline like the 2D/3D converters build_cpu_pipeline puts around
  SpatialTransform, and carry no registry entry. The feature is reached through
  pipeline.retain_stats, which smauglab/config.py validates key by key and value by
  value, and smauglab/configs/transform_params_gpu_stats-restore.json ships it.

🤖 Generated with Claude Code