fix: make the random-choose bucket usable, and seed every draw with torch

[Hendrik-code]

Stacked on #59 — last of the four correctness PRs. Deduplication starts in the next one.

The bucket wrote into the caller's batch

RandomChooseXTransformsGPU.apply_transform did out = input, then assigned back through out[i:i+1] = xi per element. It clones now, as every sibling in gpu/spatial.py already does.

The bucket could not run half the transforms it accepts

It calls apply_transform directly, which skips the forward_parameters step that fills params. Any child with a kornia _param_generator — the spatial transforms — raised:

ValueError: params must contain 'scale' tensor

It now samples those parameters itself for children that have a generator, and leaves the contrast transforms — which sample inside apply_transform — on the empty-params path they expect.

RandomLowResTransformGPU raised KeyError outside the mask path

It read flags["data_keys"] unguarded. Only MaskSequentialOpsCustom injects that key, so every other caller failed: calling the transform standalone, or from inside the bucket, which passes the transform's own flags. It defaults to IMAGE, which is what those callers mean. The MASK branch still selects nearest-neighbour resampling, and there's a test pinning that.

torch.manual_seed did not reach every draw

Blur sigmas and RandConv kernel sizes were drawn with Python's random:

• torch.manual_seed does not seed Python's random, so a "seeded" training run was not reproducible.
• Under DDP each rank has its own random state, so ranks picked different sigmas for the same batch.

The suite hid this: unit_tests/helpers.py::seed_everything seeds torch, numpy and random. Training does not call it.

New smauglab/transforms/rng.py is where those draws live now. It is built from the _next_shared_seed / _shared_rand pair that was already sitting in gpu/fromSeg.py, written for exactly this and never called once — that dead copy is deleted here rather than left as a third way to draw a random number.

Testing

unit_tests/test_bucket_and_rng.py, 12 tests. Five fail against the parent commit.

FAILED TestLowResRunsOutsideTheMaskPath::test_it_runs_standalone
FAILED TestLowResRunsOutsideTheMaskPath::test_it_runs_inside_a_random_choose_bucket
FAILED TestBucketDoesNotMutateItsInput::test_the_callers_tensor_is_left_alone
FAILED TestBucketDoesNotMutateItsInput::test_the_same_on_batch_path_also_runs_a_generator_transform
FAILED TestTorchSeedReachesEveryDraw::test_randconv_is_reproducible_under_torch_seed_alone

The RNG one had to be verified on its own, by restoring just the random.choice call site — reverting all four source files together makes the module unimportable. The seeding tests deliberately call torch.manual_seed only, never seed_everything, since seeding random is what masked the bug.

test_it_still_returns_something_transformed guards against the clone turning the bucket into a no-op.

Full suite: 69 passed, 168 subtests. ruff and mypy clean.

Compatibility

Models trained before this change saw the old behaviour and will not reproduce against it. No config key, parameter or default changed.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[NathanMolinier]

looks good