Fix division by zero in Warp for singleton spatial dimensions

[VenkateswarluNagineni]

Description

Warp.forward (native grid_sample path) normalizes grid coordinates with:

grid[..., i] = grid[..., i] * 2 / (dim - 1) - 1

When a spatial dimension has size 1 — a single-slice volume (B, C, 1, H, W) or a single-row/column image (B, C, 1, W) / (B, C, H, 1) — dim - 1 == 0, so the normalization divides by zero and produces inf/nan coordinates. With padding_mode="zeros", even a zero displacement field then yields:

• nan output for 2D single-row/column inputs, and
• all-zeros output for 3D single-slice inputs,

instead of returning the input image unchanged.

Fix

Clamp the denominator to 1 so the lone voxel maps to -1. This mirrors the guard MONAI already applies in monai.networks.utils.normalize_transform, which performs the identical align_corners=True normalization and clamps the size with norm[norm <= 1.0] = 2.0 to avoid exactly this division by zero. Non-singleton dimensions are numerically unchanged.

grid[..., i] = grid[..., i] * 2 / max(dim - 1, 1) - 1

Verification

Added test_singleton_spatial_dim (2D single-row, 2D single-column, 3D single-slice). With a zero displacement field the warped output now equals the input and contains no nan. The test fails on the current code and passes with the fix; the existing Warp tests are unaffected.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.

[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: 894527ad-b7be-4e4d-8f9c-7d0a71c771f6

📥 Commits

Reviewing files that changed from the base of the PR and between b46fe56 and a2811e2.

📒 Files selected for processing (1)

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

🚧 Files skipped from review as they are similar to previous changes (1)

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

---

📝 Walkthrough

Walkthrough

Warp.forward clamps the PyTorch fallback normalization denominator to avoid division by zero for singleton spatial dimensions. A regression test forces the non-compiled path, runs Warp on shapes with size-1 spatial dimensions, and checks for no NaNs and unchanged output with zero displacement.

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 50.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 clearly states the main fix for singleton spatial dimensions.
Description check	✅ Passed	The description follows the template and includes the change summary and checklist, with only the issue reference left blank.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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]

Actionable comments posted: 1

🧹 Nitpick comments (1)

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

141-141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a method docstring for the new test definition.

As per path instructions, "Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings."

🤖 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` at line 141, The new test definition
in test_singleton_spatial_dim is missing the required Google-style method
docstring. Add a concise docstring to this test method describing what it
verifies, including any relevant args, return value, and raised exceptions
sections if applicable, and keep the style consistent with the other test
methods in the class.

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.

Inline comments:
In `@tests/networks/blocks/warp/test_warp.py`:
- Around line 146-150: The regression test in test_warp should explicitly
exercise the native grid_sample path instead of relying on environment-dependent
fallback behavior. Update the test around Warp(...) and its result call to force
the native branch using the existing Warp/test helper mechanism or a direct flag
on Warp so the checked cases always run through grid_sample regardless of
runtime setup.

---

Nitpick comments:
In `@tests/networks/blocks/warp/test_warp.py`:
- Line 141: The new test definition in test_singleton_spatial_dim is missing the
required Google-style method docstring. Add a concise docstring to this test
method describing what it verifies, including any relevant args, return value,
and raised exceptions sections if applicable, and keep the style consistent with
the other test methods in the class.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 365796e2-2785-4f72-a3e4-26f2f4c1a718

📥 Commits

Reviewing files that changed from the base of the PR and between 557ffaa and b46fe56.

📒 Files selected for processing (2)

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

[VenkateswarluNagineni]

Thanks for the review. Addressed both points in the latest commit:

• The test now pins USE_COMPILED=False so it always runs through the native grid_sample branch — the only path that normalizes the grid (the csrc grid_pull path doesn't, so it's unaffected by this fix).
• Added a Google-style docstring to test_singleton_spatial_dim.

Verified locally: the test fails on the pre-fix code and passes with the fix; isort/black/flake8 are clean on both changed files.