Fix: handle precomputed crops and safe no-grad cleanup in auto3dseg

[bluehyena]

Fixes #8802
Also addresses #5889

Description

This pull request fixes two issues in monai.auto3dseg.analyzer.

• ImageStats.__call__() now correctly handles the optional precomputed nda_croppeds input path instead of raising UnboundLocalError.

• ImageStats, FgImageStats, and LabelStats now use an exception-safe no-grad context so analyzer failures do not leak the global PyTorch grad-enabled state, which is related to data analyzer set_grad_enabled #5889.

I also added tests covering both behaviors.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes fix a bug where ImageStats crashed when precomputed nda_croppeds was provided in the input dict, causing an UnboundLocalError. The implementation refactors autograd state management from manual toggle (torch.set_grad_enabled) to the torch.no_grad() context manager across three analyzer methods, and consolidates CUDA detection with statistics computation within this context. Tests verify precomputed nda_croppeds handling and gradient state restoration when exceptions occur.

Estimated code review effort

monai/auto3dseg/analyzer.py:
🎯 3 (Moderate) | ⏱️ ~25 minutes

tests/apps/test_auto3dseg.py:
🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% 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	Title clearly summarizes the main changes: handling precomputed crops and safe no-grad cleanup.
Linked Issues check	✅ Passed	Changes directly address #8802 (precomputed nda_croppeds handling) and #5889 (grad-state leakage on exceptions).
Out of Scope Changes check	✅ Passed	All changes are in-scope: analyzer refactoring and test coverage for the two documented fixes.
Description check	✅ Passed	PR description covers the issues fixed, implementation changes, and test additions. All required template sections are present and substantive.

_{✏️ 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 and usage tips.}

[coderabbitai]

🧹 Nitpick comments (2)

monai/auto3dseg/analyzer.py (2)

236-252: Misplaced docstring.

The docstring block is placed after the validation code. Move it to immediately follow the def __call__(self, data): line for proper function documentation.

Proposed fix

The docstring (lines 236-252) should be moved to right after line 219 (def __call__(self, data):), before the validation code.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/auto3dseg/analyzer.py` around lines 236 - 252, The function-level
docstring for __call__ is misplaced after validation; move the entire docstring
block currently referencing ImageStatsKeys, SampleOperations, and report_format
to immediately follow the def __call__(self, data): signature so it becomes the
function's official docstring (before any validation or logic), preserving the
existing text and formatting.

---

271-273: Consider strict=True for zip.

Adding strict=True catches length mismatches early if SHAPE and SPACING ever diverge unexpectedly.

-            report[ImageStatsKeys.SIZEMM] = [
-                a * b for a, b in zip(report[ImageStatsKeys.SHAPE][0], report[ImageStatsKeys.SPACING])
-            ]
+            report[ImageStatsKeys.SIZEMM] = [
+                a * b for a, b in zip(report[ImageStatsKeys.SHAPE][0], report[ImageStatsKeys.SPACING], strict=True)
+            ]

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@monai/auto3dseg/analyzer.py` around lines 271 - 273, The current computation
of ImageStatsKeys.SIZEMM uses zip(report[ImageStatsKeys.SHAPE][0],
report[ImageStatsKeys.SPACING]) which can silently truncate if lengths differ;
update the comprehension in analyzer.py to use zip(..., strict=True) so
mismatched lengths raise immediately (ensure environment uses Python 3.10+);
reference the report dict keys ImageStatsKeys.SHAPE, ImageStatsKeys.SPACING and
the target ImageStatsKeys.SIZEMM when making this change.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@monai/auto3dseg/analyzer.py`:
- Around line 236-252: The function-level docstring for __call__ is misplaced
after validation; move the entire docstring block currently referencing
ImageStatsKeys, SampleOperations, and report_format to immediately follow the
def __call__(self, data): signature so it becomes the function's official
docstring (before any validation or logic), preserving the existing text and
formatting.
- Around line 271-273: The current computation of ImageStatsKeys.SIZEMM uses
zip(report[ImageStatsKeys.SHAPE][0], report[ImageStatsKeys.SPACING]) which can
silently truncate if lengths differ; update the comprehension in analyzer.py to
use zip(..., strict=True) so mismatched lengths raise immediately (ensure
environment uses Python 3.10+); reference the report dict keys
ImageStatsKeys.SHAPE, ImageStatsKeys.SPACING and the target
ImageStatsKeys.SIZEMM when making this change.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: fb6ef620-5777-4a34-a6c5-5372ddb90b86

📥 Commits

Reviewing files that changed from the base of the PR and between 853f702 and d9a2bc4.

📒 Files selected for processing (2)

• monai/auto3dseg/analyzer.py
• tests/apps/test_auto3dseg.py

[williams145]

I noticed a couple of technical differences worth flagging between this approach and the fixes in #8809 and #8810.

On nda_croppeds handling: The ternary here has no validation on the pre-computed value , if a caller passes a wrong type or a list with the wrong number of entries, the error surfaces downstream as a confusing shape mismatch rather than a clear message. #8809 adds explicit type and length validation with a descriptive ValueError.

On grad state: torch.no_grad() always restores grad to enabled after the block, regardless of what state the caller was in before. If the caller had grad disabled before calling the analyzer, this silently re-enables it , which is a subtle semantic change. The try/finally pattern in #8810 uses torch.is_grad_enabled() to capture and restore the exact pre-call state, so both enabled→enabled and disabled→disabled are handled correctly.

[ericspod]

Hi @bluehyena thanks for looking into this. I have minor comments to make, we should also look at @williams145 's PRs though I commented in there that using no_grad should work correctly and so I'm waiting to hear back from them.

[ericspod]

This docstring comment appears spurious here, although it was in the original code please place it at the top of the method and check it's contents are correct.

[ericspod]

One other thing to note is that no_grad could be used as a method decorator as well. This might be cleaner for this code but you would have to check it is working as intended.