Fix non-functional jitter in Warp.get_reference_grid (#8953)

### Description

Little messy description as trying to fix multiple things but the jist
is: `Warp.get_reference_grid` never applied the `jitter` it advertises
and crashed whenever `jitter=True`.

The grid is built from `torch.arange` (integer dtype) and
`self.ref_grid` was assigned `grid.to(ddf)` before the jitter block, so
`grid += torch.rand_like(grid)` mutated a local that was never returned,
and `torch.rand_like` raises `NotImplementedError` on an integer tensor
anyway. Separately, `fork_rng(enabled=seed)` disabled RNG forking when
`seed` took its default of `0`, leaking the seeded state into the global
RNG.

The grid is now cast to `ddf` before jittering, the jittered tensor is
assigned to `self.ref_grid`, and `fork_rng()` isolates the seeded draw.

The non-jitter path is unchanged. A regression test covers the
float/non-integer jittered grid, the integer un-jittered grid, and
per-seed reproducibility; it fails before this change with
`NotImplementedError`.

### Types of changes
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] New tests added to cover the changes.

---------

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
Signed-off-by: Soumya Snigdha Kundu <soumyawork15@gmail.com>
Signed-off-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: aymuos15

monai/networks/blocks/warp.py

@@ -121,12 +121,13 @@ def get_reference_grid(self, ddf: torch.Tensor, jitter: bool = False, seed: int
         mesh_points = [torch.arange(0, dim) for dim in ddf.shape[2:]]
         grid = torch.stack(meshgrid_ij(*mesh_points), dim=0)  # (spatial_dims, ...)
         grid = torch.stack([grid] * ddf.shape[0], dim=0)  # (batch, spatial_dims, ...)
-        self.ref_grid = grid.to(ddf)
+        grid = grid.to(ddf)
         if jitter:
             # Define reference grid on non-integer values
-            with torch.random.fork_rng(enabled=seed):
+            with torch.random.fork_rng():
                 torch.random.manual_seed(seed)
                 grid += torch.rand_like(grid)
+        self.ref_grid = grid
         self.ref_grid.requires_grad = False
         return self.ref_grid
 

tests/networks/blocks/warp/test_warp.py

@@ -139,6 +139,21 @@ def test_ill_shape(self):
         with self.assertRaisesRegex(ValueError, ""):
             warp_layer(image=torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), ddf=torch.zeros(1, 2, 3, 3))
 
+    def test_jitter(self):
+        ddf = torch.zeros(1, 2, 4, 5)
+        grid = Warp(jitter=True).get_reference_grid(ddf, jitter=True, seed=0)
+        self.assertTrue(grid.is_floating_point())
+        self.assertFalse(torch.equal(grid, grid.round()))
+
+        grid = Warp().get_reference_grid(ddf, jitter=False)
+        self.assertTrue(torch.equal(grid, grid.round()))
+
+        same = Warp().get_reference_grid(ddf, jitter=True, seed=7)
+        repeat = Warp().get_reference_grid(ddf, jitter=True, seed=7)
+        other = Warp().get_reference_grid(ddf, jitter=True, seed=8)
+        self.assertTrue(torch.equal(same, repeat))
+        self.assertFalse(torch.equal(same, other))
+
     @mock.patch("monai.networks.blocks.warp.USE_COMPILED", False)
     def test_singleton_spatial_dim(self):
         """