[Hendrik-code]

What this is

in_seg and out_seg have been in the GPU augmentations for a long time, and between them they expressed about a third of what they looked like they expressed. This finishes them, and adds the two things they implied but never had: sub-label granularity, and a background that is not all-or-nothing.

Builds on hm/ta-mix-alpha (merge that first), and follows that PR's structure deliberately: one shared module with the rationale in its docstring, a configure_* call per constructor, and class-level "off" defaults so mypy can see the attributes.

What was broken

• in_seg and out_seg could not express a mixture. Two independent uniforms, and both fired fell through to the same "all" as neither fired — so in_seg=1, out_seg=1 was indistinguishable from in_seg=0, out_seg=0.
• mix_in_out never worked on the layout training uses. It keyed its keep/drop draw on seg_mask.shape[:2], the channel axis. For nnU-Net's single-channel integer label map that is (B, 1) — one draw per sample that deletes the whole mask half the time. So "in" applied the transform nowhere and "out" applied it everywhere. It only ever meant "per label" on a one-hot mask. Every config that set mix_in_out: true has been running a coin flip, not a label mix.
• The background was one blob. out_seg took all of it or none.
• The composite was a step function. m * transformed + (1 - m) * orig with binary m left a one-voxel discontinuity sitting exactly on the label boundary — an edge perfectly correlated with the segmentation target.
• Four of the seventeen TAs took neither parameter (RandomSynthSegGPU, RandomPaletteGPU, PaletteSynthesisGPU, RandomDomainTransferGPU) — precisely the set that already holds the segmentation.
• No validation. _choose_region_mode silently clamped, so in_seg: 5.0 and in_seg: true were accepted quietly. And it drew its two uniforms whenever a mask was attached, including at zero, so an unselected run and a selected one at the same seed were never comparable — the same defect fix: pin the region-moment matmul to fp32, which autocast overflows #136 removed for mix_prob.

The surface

in_seg: float = 0.0          # P(confine to the foreground labels)
out_seg: float = 0.0         # P(confine to the background)
                             #   both fired -> a mixture; neither -> the whole patch
seg_unit: str = "whole"      # "whole" | "label" | "component"
seg_unit_prob: float = 0.5   # P(keep) per unit
bg_split: str = "none"       # "none" | "em"
bg_em_clusters = (3, 10)     # SynthSeg's range
seg_feather: float | [lo, hi] = 0.0   # Gaussian sigma in voxels; 0 is a hard edge

Everything resolves to one integer unit id map and a single gather off it. Nothing per-unit is materialised — a noisy mask under seg_unit="component" can hold hundreds of components, and [2, 300, 128, 128, 128] of bool is 1.2 GB.

bg_split="em" calls the existing em_subdivide_labels (SynthSeg §5.4) with n_foreground_clusters=1; the EM is not reimplemented.

connected_components_3d, the one new algorithm

Label-equivalence propagation with a pointer jump, pure torch, on device, int64 throughout — float32 is exact only to 2²⁴ and a 256³ patch has exactly that many voxels. 6-connected, matching scipy.ndimage.label's default; the obvious single max_pool3d(kernel=3) is a 3³ neighbourhood and would merge corner-touching structures.

Components respect label boundaries in one propagation pass: the step is gated on same-label neighbours, so it does not need a pass per label value (a spine patch has five to twenty-five).

Checked against scipy rather than against itself — partition equality, not just counts:

case	result
random blobby volumes, 6 seeds	identical partition, every seed
multi-label volumes, 4 seeds	count equals the sum of scipy's per-label runs
two voxels touching at a corner	2 components (pins 6- over 26-connectivity)
two voxels touching at a face	1 component
serpentine ribbon	identical partition (the pointer jump's reason to exist)
empty / full / single voxel	no raise, count never 0

Measured at the real patch size

A40, 128³, batch 2, 6 labels. Full output in the verification section below.

connected_components_3d	43.7 ms / batch
background_em_units (EM, 3–10 clusters)	15.6 ms / batch
feather blur, sigma 1–2	0.9–1.1 ms / call
cache saves, 24 transforms, component + em	1443 ms / batch

That last line is why the decomposition is cached rather than recomputed per transform: without it the feature adds over a second to every training step. The feather blur is the one part that cannot be cached (it depends on the per-call draws) and at ~1 ms it does not need to be.

All 24 (mode, seg_unit, bg_split) combinations were checked at 128³ to change exactly a union of whole units, per sample — not merely "something changed". out on the whole background is 93.5% of voxels; with bg_split="em" it becomes 55.8%, and both, component, em 41.6%, which is the feature doing what it is for.

RandomRedistributeSegGPU is the one exception

Its in_seg never meant what in_seg means elsewhere — it confines each region's Gaussian to that region's own voxels — and 28 config files set it. It becomes confine_prob, on both backends, and the class deliberately does not accept in_seg / out_seg.

That is not fastidiousness. Had it accepted them, every stale config would have kept validating while silently changing meaning, and no diagnostic was possible: difflib scores in_seg against out_seg at 0.615, above unknown_parameter_message's 0.6 cutoff, so it would answer "did you mean out_seg?" and never reach the hint. Unaccepted, the hint fires:

RandomRedistributeSegGPU: unknown parameter 'in_seg'.
  'in_seg' -> on RandomRedistributeSegGPU it is now confine_prob; elsewhere in_seg is unchanged

What the class does take is the rest of the surface, which its per-region machinery consumes directly. Its defaults (seg_unit="label", seg_unit_prob=1.0) are the region set and participation it has always had, and the keep draw spends nothing at 1.0 — verified tensor-identical to the seg_region_masks set it replaces, including on a patch with no background at all.

Three stream properties, each one the old code broke

Tested, not asserted in prose:

1. A knob that is off spends no random numbers — in_seg = out_seg = 0, and seg_unit_prob = 1, leave the stream exactly where they found it. Without this no sweep is comparable with its own baseline at a fixed seed.
2. The draw width depends on the configuration, never on the data — a two-label patch and a five-label one land in the same place. A per-unit draw would spend one uniform per label present, and every transform downstream would diverge on the data.
3. The EM draw is forked — it spends ~B × 100 001 uniforms on a cache miss and none on a hit, so without the fork the augmentation would depend on cache state: on which transform asked first, on which rows p selected, on whether a geometric transform replaced the mask.

Deliberate behaviour changes

• mix_in_out is gone, migrated to seg_unit="label". This is a fix, not a rename (see above), and the draw also moves from randint(0, 2, ...) to rand(...) < prob.
• Both probabilities firing now means a mixture rather than the whole patch. Affects the three configs that set both above zero — transform_params_gpu_inoutseg.json and transform_params_hybrid{,_TAGE}.json, all at 0.5/0.5 — whose distribution goes from 25% in / 25% out / 50% all to 25 / 25 / 25% mixture / 25% all. Left as they are: the new reading is what they should have meant.
• The mode draw is per sample, not batch-wide. It was one decision for the whole batch at ten of the eleven call sites, so in_seg=0.5 read as "half the batches" while p and mix_prob both meant "half the samples". RandomInverseGPU's per-sample loop existed only to get a per-sample mode and is gone with it; its non-finite guard is now per channel, which matters only for an input that was already broken, since the transform is max - x.
• Out-of-range probabilities are rejected instead of clamped.
• RandomPaletteGPU and PaletteSynthesisGPU gain the non-finite guard the other transforms have had; they used to write the synthesis unconditionally.

Config hashes change

config_hash feeds the experiment directory name. Seven tracked configs move, and both benchmark arms named in CLAUDE.md are among them:

config	before	after
all_augmentations.json (generated)	f47a0f39	e3b5ca6a
transform_params.json	088b9bd5	8c8768b8
transform_params_gpu.json	49138fe4	c80bd868
transform_params_hybrid.json	c95704cd	c482b7d3
transform_params_hybrid_TAGE.json	c5fb06c0	726e91fd
transform_params_one-sequence-to-segment-them-all.json	c69b9872	1690ad25
transform_params_paper-synthseg.json	580beb13	5b0c6e1b

c69b9872 is the Ours (fp32) / Ours (post) arms and 580beb13 is SynthSeg (post). Existing run directories keep their names; a re-run from these files lands in new ones. The old files are in git history.

A further 28 config files in smauglab/configs/ and smauglab/configs_paul/ are gitignored local sweep configs (the blanket *.json rule). They are migrated on disk so local runs keep working, but they are not in this diff — worth knowing if you have your own.

The four synthesising TAs

Each composites after its own stats restore, which is where the region selection already sat in the other twenty-four:

• RandomSynthSegGPU, RandomPaletteGPU, PaletteSynthesisGPU with normalize=True — reviving a parameter _apply_region_mode had carried with no caller. Their output is a [0, 1] synthesis and the patch it lands on is z-scored, so a raw composite would leave an intensity step on the mask boundary that has nothing to do with the augmentation.
• RandomDomainTransferGPU after the z-restore rather than after its LUT-space blend, which is the point at which both sides share a scale.

RandomSynthSegGPU takes the parameters explicitly: _filter_generator_kwargs drops anything outside _GENERATOR_KEYS, so they would have vanished silently. Its docstring says that bg_em_clusters and the generator's own em_background_clusters_range are different features one word apart.

Implementation notes worth a reviewer's eye

The cache is a tensor attribute, not a dict keyed on tensors. Tensor.__hash__ is identity-based but Tensor.__eq__ is elementwise, so a hash-bucket collision between two masks makes dict evaluate bool(a == b) and raise "Boolean value of Tensor with more than one value is ambiguous". Measured: a WeakKeyDictionary with five hundred same-shaped tensor keys raises reliably. An attribute has no buckets, dies with the tensor, and invalidates for free — MaskSequentialOpsCustom.transform returns the same object for an image-only module and a new one for a geometric one, so identity changes exactly when the mask changes.

Keyed on the mask as the transform was handed it, which is what keeps the image and the units row-aligned: kornia slices the image to the rows that drew below p and _seg_for_applied_rows slices the mask to match, so a decomposition built from the full batch would pair B-row units with a B'-row image — and bg_split="em" reads that image. At p = 1 the slice is the full mask object and every module of the pass hits one entry. This also makes the RandomChooseXTransformsGPU bare-tensor path work with no change to it: verified 1 decomposition for two children at same_on_batch=True, and 2 for two samples × two children otherwise, rather than 4.

em_subdivide_labels is guarded against a patch with no background, where it falls back to the largest-area label and would otherwise split a foreground structure into three to ten pieces and call them background.

gpu/regions.py imports em_subdivide_labels function-locally. synthseg/__init__.py imports synthseg/transforms.py, which imports gpu/base.py, which imports this module for its defaults. A module-scope import closes that loop and breaks every import smauglab.transforms.gpu.base; test_imports.py is the canary.

The pinned-image caveat, stated rather than hidden: bg_split="em" is the only decomposition that reads the image, and every transform's first statement is input = input.clone(), so keying the cache on the image would mean a zero hit rate. The clusters are fitted once per pass, from the image as it stood when the first region-selecting transform asked. That is also the paper's reading — EM is a preprocessing step on the real paired image — and it means two region-selected transforms in one pass cut the background along the same lines instead of each along its own.

normalise_blend_range gains bounds= so seg_feather, a sigma in voxels with no upper limit, shares it with the three blend weights that need the unit interval. The default keeps the [0, 1] wording its existing callers' tests assert on, and a test checks that widening one did not widen mix_alpha.

Verification

pytest unit_tests -q                 556 passed, 2701 subtests, 18 skipped
                                     11 failed -- the documented seg11 baseline
                                     (9 from nnunetv2 2.4.2, 2 from the stale
                                      random-normalization config); identical
                                      before and after this branch
smauglab template --check            passes
smauglab matrix --check              passes
ruff check / ruff format --check     clean
mypy smauglab                        no issues (2.1.0 locally; CI's 2.4.0 is the real check)

Beyond the suite, on an A40 at 128³ batch 2:

=== every (mode, seg_unit, bg_split) changes exactly a union of whole units ===
  in    whole      bg=none  OK   6.5% of voxels
  in    label      bg=none  OK   4.1% of voxels
  in    component  bg=em    OK   2.8% of voxels
  out   whole      bg=none  OK  93.5% of voxels
  out   whole      bg=em    OK  55.8% of voxels
  out   label      bg=em    OK  35.6% of voxels
  both  whole      bg=none  OK 100.0% of voxels
  both  component  bg=em    OK  41.6% of voxels
  all   *          *        OK  whole patch
  -> all exact  (24 combinations)

=== the feather is a bounded, monotone ramp ===
  sigma= 0.0  ramp voxels=        0  unit core min=1.00000  far bg max=0.00e+00  bounded=True
  sigma= 1.0  ramp voxels=   355300  unit core min=1.00000  far bg max=0.00e+00  bounded=True
  sigma= 4.0  ramp voxels=  1510066  unit core min=0.98987  far bg max=0.00e+00  bounded=True
  profile out of a label along one axis is non-increasing: True

=== a six-transform pipeline with the surface on, all three container modes ===
  sequential       340.3 ms   decompositions=4   finite=True   mask untouched=True
  random_order     134.4 ms   decompositions=4   finite=True   mask untouched=True
  random_order_ta  109.8 ms   decompositions=4   finite=True   mask untouched=True
  (4 distinct specs across 6 transforms -> one decomposition each, shared)

RandomDomainTransferGPU could not be exercised locally: it needs the SMAUGLAB_DOMAIN_BANK artefact, which is not set here, so its tests skip as they do upstream. Its region selection is the simplest of the four (no normalize, both sides already on one scale) but it is the one path in this PR that has not run.

🤖 Generated with Claude Code