feat: PALETTE, composed from swappable partitioning blocks

[Hendrik-code]

Supersedes #45, ported onto the registry stack (#67 → #71). #45 is left as it is until this one is reviewed.

What this adds

RandomPaletteGPU hard-codes one composition — 1-D K-means, then Voronoi, then a per-region affine remap. PaletteSynthesisGPU takes that apart:

slot	blocks
initial_partitioner	kmeans1d, em_gmm
refinement_partitioners (a list, applied in order)	voronoi, em_gmm, identity
overlay	the per-anatomical-label remap, with a new blend_strength knob

The remap itself stays fixed — it is the part the augmentation is named after. At its defaults the class is RandomPaletteGPU; the point is the compositions that are not, such as EM/GMM instead of K-means, or several refinements in a row.

"PaletteSynthesisGPU": {
    "p": 0.25,
    "initial_partitioner": {"type": "em_gmm", "n_foreground_clusters": 3},
    "refinement_partitioners": [{"type": "voronoi", "skip_prob": 0.4}],
    "overlay": {"enabled": true, "blend_strength": [0.6, 1.0]}
}

RandomPaletteGPU is left exactly as it is. It is the composition every shipped config names, and rewriting it in terms of these blocks would put a refactor in the way of reading what the published experiments ran.

Why this is a port and not a rebase

#45 branched before the auglab → smauglab rename and edited the if ladder in gpu/transforms.py, which #70 deleted. The rebase touched every hook-up point:

• the class registers itself instead of being appended by the ladder, and gets its own AugId (palette_composed) rather than sharing PALETTE — two entries on one concept silently overwrite each other's cell in matrix();
• build_palette_from_cfg is gone with the ladder that called it. build.py builds a transform by calling its constructor with the config block, so the constructor now takes the nested blocks directly — dicts from a config, instances from a test. factory.py keeps the part the registry cannot do: signature-derived validation cannot look inside a block whose shape depends on its own type field;
• every parameter is named and everything has a default, because the registry rejects **kwargs (test_no_registered_class_hides_parameters_behind_kwargs) and requires a GPU entry to be constructible with none (test_constructible_with_declared_defaults);
• random.choice for the blur sigmas becomes shared_choice — the fix from transforms/rng.py, since torch.manual_seed does not reach Python's random and DDP ranks would each pick their own sigma.

Two bugs from #45 fixed along the way:

• the overlay's entry point is remap_labels, not apply. nn.Module.apply already means "run this over every submodule", so .apply(fn) was broken on any pipeline containing one — mypy caught it.
• the EM blocks return parent * mult + assign ids, which are sparse; the remap indexes length-R tensors with them. _densify_region_ids makes them contiguous, and a test pins it for every block.

Migration

PaletteSynthesisTransform stays in DEAD_KEYS: the class only ever existed on the palette-refactor branch, so a config carrying that key never ran it, and turning it on during migration would start running an augmentation the config never ran. The dropped-key note now names PaletteSynthesisGPU and says what to rename. This also keeps legacy_effective_kwargs.json honest as captured evidence that the registry builds what the ladders built.

Tests

unit_tests/test_palette_blocks.py covers what is new — that a config picks the blocks, that a nonsensical pick fails at construction with the fix in the message (voronoi in the initial slot, an unknown type, a block without a type), that region ids stay contiguous, and that every composition runs end to end. PaletteSynthesisGPU also joins the introspective discovery in test_transforms_gpu.py.

234 passed, 3 skipped locally, plus ruff check, ruff format --check and mypy smauglab clean. README matrix and all_augmentations.json regenerated.

🤖 Generated with Claude Code