[Hendrik-code]

The bug

em_subdivide_labels (SynthSeg EM label completion) takes each label's voxel count from a histogram of label_map.clamp_min(0), then selects that label's voxels with label_map == p.

A negative value is not a label: nnU-Net writes -1 for voxels outside the annotated region. The clamp counted every -1 as background; the selection that follows excludes them. So the count exceeded the number of voxels gathered, and the randint(0, count) subsample that fits the background GMM could draw an index past the end of x:

cnt = int(counts_per_sample[b][p])     # includes the -1 voxels
mask = lab_b == p                      # excludes them
x = ref_b[mask]                        # len(x) < cnt
sel = torch.randint(0, cnt, (max_fit_voxels,))
x_fit = x[sel]                         # out of bounds

On CPU that raises IndexError. On CUDA it is an asynchronous index out of bounds device-side assert, reported at whichever kernel happens to run next — in my case inside _em_gmm_1d, nine frames from the actual fault. It only fires when an oversized index is actually drawn, so it is probabilistic: for a 128³ patch with six ignored voxels and the shipped em_max_fit_voxels: 100000, about two runs in five.

Reproduced on a preprocessed Dataset804 patch: the histogram reports 1219605 background voxels where label == 0 selects 1219599, the difference being exactly the six ignored voxels.

Scope

No published training is affected. nnUNetTrainerDAExt.get_training_transforms puts RemoveLabelTansform(-1, 0) in the CPU dataloader pipeline, ahead of the GPU augmentations, so a GPU transform is never handed a negative label in training.

It is reachable from RandomSynthSegGPU and from em_subdivide_labels directly, both of which take a label map and document no such precondition. I hit it in an offline harness that reads the preprocessed .npy files straight off disk.

Introduced in 20360d8 ("perf: remove the per-sample loops and host stalls from the augmentations"), which replaced the correct cnt = int(mask.sum()) with the batched histogram.

The fix

1. _label_counts_per_sample excludes negative values from the histogram, keeping the single-bincount shape that commit was after. The clamp stays on the row offset only, so a negative label cannot land in another sample's bins.
2. em_subdivide_labels returns one undivided region when no non-negative label is present at all. That case was previously unreachable — the clamp turned an all-ignore patch into all-background — and would otherwise index an empty parents.

Behaviour is unchanged for any label map without negative values, which is every label map the trainer produces.

Tests

New unit_tests/test_synthseg_em_ignore_label.py, four tests. The first pins the invariant the crash violated (counted == selected) and is deterministic; the others exercise the call that raised, and the all-ignore patch.

Verified both directions on this branch:

• with the fix: 4 passed, 4 subtests passed
• with functional.py reverted: 5 failed, 1 passed — IndexError: index 2272 is out of bounds for dimension 0 with size 1840

Full suite: 473 passed, 13 skipped, 1523 subtests passed.

🤖 Generated with Claude Code