Preserve input dtype in per-instance FFT/bias transforms; guard subsample target

Ghosting, Spike, Motion and BiasField computed in float32 and only cast
back to the input dtype for applied elements, so per-instance gating
could mix dtypes and break torch.cat for non-float32 batches. Cast every
per-element output back to the input dtype (matching the Blur
convention). Added a dtype-preservation regression test.

Also reject non-positive subsample targets in _subsample_for_quantile
with a clear ValueError now that subsample_size is caller-exposed.

Author: fepegar

src/torchio/transforms/intensity/bias_field.py

@@ -236,7 +236,8 @@ def _apply_bias_per_element(
             seed=seed_per_element[index],
             device=data.device,
         )
-        outputs.append(slice_b / field if divide else slice_b * field)
+        corrected = slice_b / field if divide else slice_b * field
+        outputs.append(corrected.to(data.dtype))
     return torch.cat(outputs, dim=0)
 
 

src/torchio/transforms/intensity/ghosting.py

@@ -206,4 +206,4 @@ def _add_ghosting(
         torch.fft.ifftshift(corrupted, dim=(-3, -2, -1)),
         dim=(-3, -2, -1),
     ).real
-    return result
+    return result.to(data.dtype)

src/torchio/transforms/intensity/motion.py

@@ -176,7 +176,7 @@ def _apply_motion(
         reconstructed = torch.fft.ifftn(spectrum, dim=(-3, -2, -1))
         result[b] = reconstructed.real
 
-    return result
+    return result.to(data.dtype)
 
 
 def _apply_rigid_transform(

src/torchio/transforms/intensity/normalize.py

@@ -384,12 +384,18 @@ def _subsample_for_quantile(values: Tensor, target: int) -> Tensor:
 
     Args:
         values: 1D tensor of values.
-        target: Maximum number of elements to keep.
+        target: Maximum number of elements to keep. Must be positive.
 
     Returns:
         *values* unchanged if already within *target*, else a strided
         view with at most *target* elements.
+
+    Raises:
+        ValueError: If *target* is not a positive integer.
     """
+    if target <= 0:
+        msg = f"target must be a positive integer, got {target}"
+        raise ValueError(msg)
     if values.numel() <= target:
         return values
     step = -(-values.numel() // target)

src/torchio/transforms/intensity/spike.py

@@ -166,4 +166,4 @@ def _add_spikes(
         torch.fft.ifftshift(spectrum, dim=(-3, -2, -1)),
         dim=(-3, -2, -1),
     ).real
-    return result
+    return result.to(data.dtype)

tests/test_normalize.py

@@ -282,6 +282,17 @@ def test_subsample_for_quantile_passthrough_when_small(self) -> None:
         sample = _subsample_for_quantile(values, 1000)
         assert sample is values
 
+    @pytest.mark.parametrize("target", [0, -1])
+    def test_subsample_for_quantile_rejects_non_positive_target(
+        self,
+        target: int,
+    ) -> None:
+        from torchio.transforms.intensity.normalize import _subsample_for_quantile
+
+        values = torch.arange(10, dtype=torch.float32)
+        with pytest.raises(ValueError, match="positive integer"):
+            _subsample_for_quantile(values, target)
+
     def test_rescale_intensity_large_image(self) -> None:
         # Full integration on a tensor exceeding torch.quantile's 2**24 limit.
         # The default 0/100 percentiles use min/max, so this stays fast.

tests/test_per_instance.py

@@ -129,3 +129,41 @@ def test_history_slice_out_of_range_raises(self) -> None:
         history = result.applied_transforms
         with pytest.raises(IndexError, match="batch of size 4"):
             _slice_history(history, 4)
+
+
+class TestPerInstanceDtypePreservation:
+    """Per-instance gating must not produce mixed-dtype batch outputs.
+
+    Float-domain intensity transforms compute in ``float32`` internally.
+    When per-element gating skips some elements, the skipped ones keep the
+    input dtype while applied ones must be cast back, so ``torch.cat`` over
+    the batch does not fail on a dtype mismatch for non-``float32`` inputs.
+    """
+
+    @staticmethod
+    def _mixed_dtype_batch(dtype: torch.dtype, batch_size: int = 8):
+        data = (torch.rand(1, 8, 8, 8) + 0.5).to(dtype)
+        subjects = [
+            tio.Subject(t1=tio.ScalarImage(data.clone())) for _ in range(batch_size)
+        ]
+        return tio.SubjectsBatch.from_subjects(subjects)
+
+    @pytest.mark.parametrize(
+        "transform",
+        [
+            tio.Ghosting(num_ghosts=4, intensity=1.0, p=0.5),
+            tio.Spike(num_spikes=2, intensity=1.0, p=0.5),
+            tio.Motion(degrees=10.0, translation=10.0, num_transforms=2, p=0.5),
+        ],
+    )
+    def test_fft_transforms_preserve_float64(self, transform) -> None:
+        torch.manual_seed(0)
+        batch = self._mixed_dtype_batch(torch.float64)
+        result = transform(batch)
+        assert result.t1.data.dtype == torch.float64
+
+    def test_bias_field_preserves_float16(self) -> None:
+        torch.manual_seed(0)
+        batch = self._mixed_dtype_batch(torch.float16)
+        result = tio.BiasField(std=0.5, p=0.5)(batch)
+        assert result.t1.data.dtype == torch.float16