fix(transforms): make Crop.compute_slices torch.compile-friendly#8960
fix(transforms): make Crop.compute_slices torch.compile-friendly#8960aymuos15 wants to merge 7 commits into
Conversation
CenterSpatialCrop (and other Crop subclasses) failed under torch.compile because compute_slices round-tripped integer ROI specs through CPU tensors via convert_to_tensor(..., device="cpu"). Under tracing the shape-derived values are fake tensors whose storage is not allocated, so torch.as_tensor raised "data is not allocated yet". The ROI arithmetic is plain integer math, so compute it directly in Python. This is traceable and drops two CPU allocations per call. Behaviour is unchanged for int, sequence, tensor and ndarray ROI inputs. Fixes Project-MONAI#8191 Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds integer ROI coercion and broadcasting helpers in Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tests/transforms/test_center_spatial_crop.py (1)
57-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a docstring to
test_torch_compile.The new test definition is missing the Google-style docstring required for new Python definitions.
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/transforms/test_center_spatial_crop.py` at line 57, Add a Google-style docstring to the new test_torch_compile method in test_center_spatial_crop so it follows the project’s Python definition requirements. Update the test_torch_compile definition to include a concise docstring that describes the test purpose and, if applicable, documents arguments, return value, and exceptions in Google style to match the surrounding test conventions.Source: Path instructions
monai/transforms/croppad/array.py (1)
345-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument
_to_int_listin Google style.The new helper's one-line docstring omits
Args,Returns, and coercion/exception behavior, which makes this conversion path harder to audit.
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 `@monai/transforms/croppad/array.py` around lines 345 - 346, The helper `_to_int_list` needs a full Google-style docstring instead of only a one-line summary. Update the docstring to include an `Args` section describing `data` and its accepted input types, a `Returns` section explaining it returns a list of Python ints, and a `Raises` section covering any coercion or conversion failures. Keep the documentation on the `_to_int_list` function aligned with the style used elsewhere in `array.py`.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 `@monai/transforms/croppad/array.py`:
- Around line 397-409: The ROI slice construction in the cropping logic is
silently truncating mixed-arity inputs because the final zip over starts/ends
does not broadcast lengths consistently. Update the ROI handling in the relevant
crop transform path to match the intended broadcasting semantics used by the
roi_center/roi_size branch, and make non-singleton length mismatches fail
explicitly instead of dropping dimensions. Use the existing ROI parsing logic in
the crop/pad transform method that builds the slice tuple to keep scalar and
sequence inputs aligned, and add a regression test for mixed scalar/sequence
roi_start/roi_end and incompatible roi_center/roi_size cases.
In `@tests/transforms/test_center_spatial_crop.py`:
- Around line 58-65: The test currently hard-resets the global meta-tracking
flag in the CenterSpatialCrop compiled-path test, which can make the suite
order-dependent. Update the test around set_track_meta so it saves the previous
tracking state before disabling it, then restores that original value in the
finally block instead of always forcing True; keep the torch._dynamo.reset()
cleanup as-is.
---
Nitpick comments:
In `@monai/transforms/croppad/array.py`:
- Around line 345-346: The helper `_to_int_list` needs a full Google-style
docstring instead of only a one-line summary. Update the docstring to include an
`Args` section describing `data` and its accepted input types, a `Returns`
section explaining it returns a list of Python ints, and a `Raises` section
covering any coercion or conversion failures. Keep the documentation on the
`_to_int_list` function aligned with the style used elsewhere in `array.py`.
In `@tests/transforms/test_center_spatial_crop.py`:
- Line 57: Add a Google-style docstring to the new test_torch_compile method in
test_center_spatial_crop so it follows the project’s Python definition
requirements. Update the test_torch_compile definition to include a concise
docstring that describes the test purpose and, if applicable, documents
arguments, return value, and exceptions in Google style to match the surrounding
test conventions.
🪄 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: b22e25bb-8d73-4264-ad2e-3ea4d2771f5e
📒 Files selected for processing (2)
monai/transforms/croppad/array.pytests/transforms/test_center_spatial_crop.py
The pure-Python rewrite zipped ROI lists directly, which silently truncated mismatched-length inputs and dropped scalar broadcasting for roi_start/roi_end. Restore the original semantics: broadcast a scalar against a sequence and raise on incompatible non-scalar lengths, via a shared _broadcast_int_pair helper. Adds a regression test for mixed scalar/sequence ROI inputs. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/transforms/test_center_spatial_crop.py (1)
57-57: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Google-style docstrings to these test methods.
Both modified definitions are missing docstrings, which breaks the repo rule for Python definitions. As per path instructions, "
**/*.py: ... 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."Also applies to: 66-66
🤖 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/transforms/test_center_spatial_crop.py` at line 57, Add Google-style docstrings to the missing Python test methods in this test class, especially test_compute_slices_broadcast and the other modified test definition referenced in the review. Update each method to include a brief summary and any appropriate Args/Returns/Raises sections if applicable, matching the repo’s docstring rule for all Python definitions.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 `@monai/transforms/croppad/array.py`:
- Around line 357-359: The ROI-to-list conversion in the helper around the
`Sequence` check is too permissive because `str` and `bytes` are treated as
sequences, causing values like "10" to be split into characters and converted
incorrectly. Update this conversion logic to reject string/byte inputs before
the `Sequence` branch, and only accept true numeric sequences or scalar numeric
values so the caller fails fast instead of building wrong slices. Use the
existing conversion helper in `array.py` to keep the validation consistent with
the rest of the cropping/padding code.
---
Nitpick comments:
In `@tests/transforms/test_center_spatial_crop.py`:
- Line 57: Add Google-style docstrings to the missing Python test methods in
this test class, especially test_compute_slices_broadcast and the other modified
test definition referenced in the review. Update each method to include a brief
summary and any appropriate Args/Returns/Raises sections if applicable, matching
the repo’s docstring rule for all Python definitions.
🪄 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: 6917eebb-952f-4463-aeca-5659afe6f222
📒 Files selected for processing (2)
monai/transforms/croppad/array.pytests/transforms/test_center_spatial_crop.py
set_track_meta is a process-global flag; save and restore the prior value instead of forcing it back to True, so the test is not order-dependent. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
|
On the remaining nitpicks: MONAI test files follow a no-docstring convention. B905 is in the repo's ruff ignore list in pyproject.toml. |
str/bytes are Sequences, so a value like "10" would be silently parsed into [1, 0]; raise TypeError instead and cover it with a test. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
Trim the oversized docstrings on the private _to_int_list/_broadcast_int_pair helpers to one-line summaries, and unify the end-clamp so it lives in a single place in the final slice comprehension instead of being re-derived in both branches. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
monai/transforms/croppad/array.py (1)
345-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse Google-style docstrings for these helpers.
These one-liners do not document args, returns, or raised exceptions, which is required for new Python definitions here. 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."
Also applies to: 356-359
🤖 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/transforms/croppad/array.py` around lines 345 - 346, The helper docstrings in _to_int_list and the other affected definitions should be updated to Google-style format. Add proper Args, Returns, and where applicable Raises sections for each new Python definition in this area, describing every parameter and the return value clearly so the helpers are fully documented.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.
Nitpick comments:
In `@monai/transforms/croppad/array.py`:
- Around line 345-346: The helper docstrings in _to_int_list and the other
affected definitions should be updated to Google-style format. Add proper Args,
Returns, and where applicable Raises sections for each new Python definition in
this area, describing every parameter and the return value clearly so the
helpers are fully documented.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 37dfbd5a-2a74-4706-bc63-6364fb361faf
📒 Files selected for processing (1)
monai/transforms/croppad/array.py
The Inductor backend forces a real C++ compilation, which fails on CI runners without a working toolchain (no MSVC on Windows) or with a clearml-corrupted sys.path on Linux. The eager backend still traces the transform through Dynamo, which is what this test verifies. Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
1 participant
CenterSpatialCropblows up undertorch.compilewhile the other crop transforms are fine (#8191). It fails inCrop.compute_sliceswithThe tensor has a non-zero number of elements, but its data is not allocated yet.The cause is that
compute_slicesran its start/end math through CPU tensors (convert_to_tensor(..., device="cpu")). ForCenterSpatialCropthe ROI values come from the input shape, so under tracing they're fake tensors with no storage, and moving them to the CPU asks Dynamo for data that isn't there.Since it's just integer math, I moved it to plain Python. A small
_to_int_listhelper handles the input forms (scalar, sequence, tensor, ndarray), with the same clamping and broadcasting as before.CenterSpatialCropnow compiles like the rest of the transforms.Added a regression test that compiles
CenterSpatialCropand checks the shape (fails before, passes after), guarded for PyTorch versions withtorch.compile.Fixes #8191.