fix: make the spatial transforms sample the way they claim to

[Hendrik-code]

Three defects in gpu/spatial.py, all silent — nothing crashes and no test fails, the pipeline just does something other than what the config asked for.

RandomFlipTransformGPU never read the flip flags its own generator sampled

The loop recomputed the same flip_axis-derived list for every batch element, so it flipped all configured axes, identically, on every call. Three seeded calls gave byte-identical output, and FlipGenerator3D — including its "at least one axis" guarantee — was dead code. It now reads params["flip"], falling back to the old all-configured-axes behaviour only for callers that reach into apply_transform directly with no sampled params.

The single-axis generators drew their "random" axis in make_samplers

kornia calls make_samplers once and caches the samplers it builds, so the axis was fixed for the transform's lifetime — "degrade a random axis" degraded the same axis for a whole training run. The draw moves to forward(), via new _choose_axis / _keep_one_axis helpers shared by ScaleGenerator3D and CropGenerator3D.

CropGenerator3D placed the crop and its position on different axes

It drew a separate dim for each, so the crop could be taken along one axis while the position placing it was randomised along another. It also neutralised the position to 1.0 by copying the crop's neutral value; the position is the crop centre as a fraction of the axis, so its neutral value is 0.5 (centred), not the far edge.

Also, one hazard rather than a fix

apply_transform_mask used "resample" in flags plus a bare assignment: resample_method was annotated but only assigned inside the if, so the restore below it could read an unbound local. It is unreachable today — apply_transform indexes flags["resample"] and raises KeyError first — so flags.get(...) removes a hazard rather than fixing an observed failure. Its two tests in TestMaskResampleRestore are guards and pass either way; the file's docstring says so.

Testing

unit_tests/test_spatial_sampling.py, 12 tests. Five of them fail against the parent commit, verified by stashing the source change and re-running:

FAILED TestFlipIsActuallyRandom::test_batch_elements_flip_independently
FAILED TestFlipIsActuallyRandom::test_two_seeds_give_two_different_flips
FAILED TestSingleAxisIsRedrawnEveryCall::test_the_crop_generator_neutralises_position_at_the_centre
FAILED TestSingleAxisIsRedrawnEveryCall::test_the_crop_generator_uses_one_axis_for_crop_and_position
FAILED TestSingleAxisIsRedrawnEveryCall::test_the_scale_generator_does_not_pin_one_axis_forever

The rest pin behaviour that must survive the change (mask and image flip together, only configured axes are flipped, one_dim=False still leaves every axis independent).

Full suite: 25 passed, 142 subtests. ruff check and ruff format --check clean, mypy smauglab/ clean.

Compatibility

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

---

Second commit: an unrelated CI fix

dcd133d fixes the red mypy job. It is not caused by the change above — the job broke on a dependency released since main's last lint run (green on 2026-08-10):

numpy/__init__.pyi:737: error: Type statement is only supported in Python 3.12 and greater  [syntax]
Found 1 error in 1 file (errors prevented further checking)

[tool.mypy] sets python_version = "3.10" deliberately — the floor of what requires-python promises — but the job ran on 3.12. mypy applies its target version when parsing every file, third-party stubs included, so pip resolving numpy 2.5 (which requires Python >=3.12 and uses PEP 695 type statements) made mypy reject numpy's own stub. Nothing in smauglab/ was reached; the run aborted at the parse error.

Running the job on 3.10 makes interpreter and target agree — which the config comment already said was the intent — and resolves numpy 2.2.x, whose stubs parse under the target. The test (python 3.10) job installs this same dependency set already, so the install path is exercised.

Not reproducible locally (numpy 2.5 cannot be installed below 3.12). I verified instead that numpy 2.2.6 type-checks clean against python_version = "3.10", and that numpy 2.3.4 does not yet carry the offending syntax.

Happy to split this into its own PR if you'd rather review it separately — it's here because the branch is red without it.

[Copilot]

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

[NathanMolinier]

Looks ok