[Hendrik-code]

What

A new pipeline key, shuffle, that redraws the order of a transform group once per batch:

"pipeline": {"shuffle": ["TA"]}        // transfer augmentations
"pipeline": {"shuffle": ["GE"]}        // general enhancements
"pipeline": {"shuffle": ["TA", "GE"]}  // both, each among its own positions

The members of a group are permuted among the positions that group already holds in the sequential pipeline. Everything else stays where it is: geometry in front, the other group's transforms between them, force_sequential transforms (LowRes), and each transform's own p and batch semantics. The only transforms that move the mask (Flip, Affine, Crop, LowRes) are never reordered, so the image/mask pairing is untouched.

Why not the existing random_order modes

They also reorder, but they change more than order, so they can't answer "does the order matter":

• they pull each group into a RandomChooseXTransformsGPU bucket, which puts every TA ahead of every GE. On the paper config that moves the five elementwise functions and the bias field from after the GE block to before it, and hoists LowRes ahead of all TAs;
• the bucket calls apply_transform directly instead of going through kornia's forward.

A shuffled sequential pipeline differs from the unshuffled one in the order alone. Without shuffle, or with shuffle: [], it builds exactly the pipeline it builds today (tested: byte-identical output under one seed).

How

• SmaugConfig.shuffle_groups() reads the key. Validation rejects unknown or duplicate groups, GEO, and shuffle combined with a bucket mode.
• build.shuffle_slots(pipeline, groups) reads each group's positions back from the registry.
• AugmentationSequentialCustom(shuffle_slots=...) overrides get_forward_sequence: when no params are given it permutes each slot group with one torch.randperm. The drawn order goes into the ParamItem list, so the mask pass and any replay follow it.
• Both the nnU-Net trainer and AugTransformsGPU pass it through. The trainer's config line now ends in , shuffle: TA (etc.) when it is set.

One consequence to know about: RandomSynthSegGPU, RandomPaletteGPU and RandomDomainTransferGPU are in the TA group. Under shuffle: ["TA"] they can land behind other TAs, and since they replace the image, whatever ran before them is overwritten (SynthSeg's EM clustering still sees it). None of the shipped configs combine these with the shuffle.

Tests

unit_tests/test_pipeline_shuffle.py: config parsing and rejection, the slot set on the paper config, non-group positions fixed across draws, all 3! permutations reached, the applied order (via record_applications) matching the drawn order with the mask unchanged, and the empty-shuffle regression. The full CPU suite passes: 476 passed, 10 skipped (466 on main + 10 new). Ruff clean (local 0.17.1; CI pins 0.16.1). mypy was not run locally.

🤖 Generated with Claude Code