Fix non-functional jitter in Warp.get_reference_grid

[aymuos15]

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

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d6655d4d-1964-4b0d-91c2-8aa6210bab52

📥 Commits

Reviewing files that changed from the base of the PR and between 764f520 and ffbd7a1.

📒 Files selected for processing (1)

• tests/networks/blocks/warp/test_warp.py

✅ Files skipped from review due to trivial changes (1)

• tests/networks/blocks/warp/test_warp.py

---

📝 Walkthrough

Walkthrough

Warp.get_reference_grid now aligns the reference grid with the displacement field before applying jitter and always forks RNG state before seeding. Tests cover jittered and non-jittered grids plus seeded determinism.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately summarizes the main fix to Warp.get_reference_grid jitter behavior.
Description check	✅ Passed	The PR description matches the template well and includes a clear summary plus the required types-of-changes section.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

monai/networks/blocks/warp.py (1)

115-120: 🎯 Functional Correctness | 🟠 Major

Cache ignores jitter/seed parameters.

The early return in get_reference_grid (lines 115-120) returns the cached self.ref_grid solely based on shape. If Warp instance is reused with different jitter or seed settings, stale grids are returned.

Internal call at line 152 passes jitter=self.jitter but ignores seed. The cache key lacks these parameters, rendering them ineffective on subsequent calls with mismatched settings.

        if (
            self.ref_grid is not None
            and self.ref_grid.shape[0] == ddf.shape[0]
            and self.ref_grid.shape[1:] == ddf.shape[2:]
        ):
            return self.ref_grid  # type: ignore

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/networks/blocks/warp.py` around lines 115 - 120, The cache in
get_reference_grid only checks tensor shape, so it can return a stale
self.ref_grid when Warp is reused with different jitter or seed values. Update
the cache validation so it also accounts for the jitter/seed inputs used to
build the grid, and make sure the call site in Warp that passes self.jitter also
propagates seed consistently. If the stored grid was created with different
settings, regenerate it instead of returning the cached value.

🧹 Nitpick comments (1)

tests/networks/blocks/warp/test_warp.py (1)

141-154: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Good coverage. Float dtype, non-integer jitter, integer-aligned baseline, and seed determinism all checked.

Optional: add a Google-style docstring; path instructions ask for docstrings on all definitions.

As per path instructions: "Docstrings should be present for all definition".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/networks/blocks/warp/test_warp.py` around lines 141 - 154, The test
coverage in test_jitter is solid, but the new path-level requirement says every
definition should have a Google-style docstring. Add a concise docstring to the
test_jitter method in the Warp test class, describing the jitter behavior it
verifies and keeping it consistent with the existing test naming and structure.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@monai/networks/blocks/warp.py`:
- Around line 115-120: The cache in get_reference_grid only checks tensor shape,
so it can return a stale self.ref_grid when Warp is reused with different jitter
or seed values. Update the cache validation so it also accounts for the
jitter/seed inputs used to build the grid, and make sure the call site in Warp
that passes self.jitter also propagates seed consistently. If the stored grid
was created with different settings, regenerate it instead of returning the
cached value.

---

Nitpick comments:
In `@tests/networks/blocks/warp/test_warp.py`:
- Around line 141-154: The test coverage in test_jitter is solid, but the new
path-level requirement says every definition should have a Google-style
docstring. Add a concise docstring to the test_jitter method in the Warp test
class, describing the jitter behavior it verifies and keeping it consistent with
the existing test naming and structure.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5f4e987b-6ff9-4a26-b3f3-eafcff5b7b94

📥 Commits

Reviewing files that changed from the base of the PR and between 4ba89bd and e3d9aeb.

📒 Files selected for processing (2)

• monai/networks/blocks/warp.py
• tests/networks/blocks/warp/test_warp.py

[ericspod]

Hi @aymuos15 super minor changes you can do if you feel needed but good to go as it is.

[aymuos15]

Thanks for pointing those out. Will remember it for the next one!