Fix channel-first indices buffer for distance_transform_edt (return_indices=True) #8656

[alexanderjaus]

Fixes #8656

Distance_transform_edt indices preallocation to use channel-first (C, spatial_dims, ...) layout for both torch/cuCIM and NumPy/SciPy paths, resolving “indices array has wrong shape” errors when return_indices=True.

Description

  import torch
  from monai.transforms.utils import distance_transform_edt

  img = torch.tensor([[[0, 0, 1],
                       [0, 1, 1],
                       [1, 1, 1]]], dtype=torch.float32)  # shape (1, 3, 3)

  # Previously raised: RuntimeError: indices array has wrong shape
  indices = distance_transform_edt(img, return_distances=False, return_indices=True)
  print(indices.shape)  # now: (1, 2, 3, 3)

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

This change fixes a shape mismatch bug in the distance_transform_edt function by correcting how the indices buffer is initialized when return_indices=True. The indices tensor is now created with a channel-first layout (C, spatial_dims...) instead of the incorrect (ndim, C, spatial_dims...) format. The fix applies uniformly to both CuPy and NumPy code paths in monai/transforms/utils.py.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Single file affected with localized, homogeneous changes
• Shape calculation logic is straightforward
• Applied consistently across CuPy and NumPy paths

Areas requiring attention:

• Verify the shape calculation correctly produces (C, D1, D2, ...) output for various tensor dimensions (2D, 3D, 4D, etc.)
• Confirm the fix resolves the SciPy/cuCIM shape requirement for the indices array
• Validate that both code paths (CuPy and NumPy) produce identical shapes with the new initialization

Pre-merge checks and finishing touches

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly references the specific bug fix (#8656) and accurately describes the channel-first indices buffer correction for distance_transform_edt.
Description check	✅ Passed	Description includes issue reference, clear explanation of the fix, reproducible code example, and completed checklist items matching the template structure.
Linked Issues check	✅ Passed	Code changes directly address issue #8656 by fixing the indices buffer allocation from (ndim, C, ...) to channel-first (C, spatial_dims, ...) layout for both CuPy and NumPy paths.
Out of Scope Changes check	✅ Passed	All changes in monai/transforms/utils.py are scoped to fixing the indices buffer allocation for distance_transform_edt, with no extraneous modifications.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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]

Actionable comments posted: 0

Caution

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

⚠️ Outside diff range comments (2)

monai/transforms/utils.py (2)

2499-2507: CuCIM path: allocate indices on img.device.

Line 2501 allocates indices without device=img.device; when img is CUDA, convert_to_cupy(indices) will fail because the CPU tensor cannot be converted to a CuPy array.

-                indices = torch.zeros((img.shape[0],) + (img.dim() - 1,) + img.shape[1:], dtype=dtype)  # type: ignore
+                indices = torch.zeros(
+                    (img.shape[0], img.dim() - 1, *img.shape[1:]),
+                    dtype=dtype,
+                    device=img.device,
+                )  # type: ignore

---

2484-2486: User-supplied output buffers are discarded by resetting to None immediately after capture.

Line 2484 captures distances and indices parameters into distances_original and indices_original, but line 2485 immediately sets them to None. This makes all downstream validation branches (lines 2492-2498, 2501-2507 in GPU path; 2514-2520, 2522-2525 in CPU path) unreachable dead code, and distances_original / indices_original are never used. The function ignores user-supplied buffers entirely.

Fix line 2485:

-    distances, indices = None, None
+    distances, indices = distances_original, indices_original

Additionally, the device/type guards at lines 2493 and 2503 use and where or is intended—they should reject if missing the tensor type or the wrong device, not requiring both conditions.

🧹 Nitpick comments (1)

monai/transforms/utils.py (1)

2533-2536: Tuple unpacking improves readability here.

The proposed refactoring is correct. scipy.ndimage.distance_transform_edt with return_indices=True expects indices with dtype int32 and shape (input.ndim,) + input.shape. Since each channel is processed as a 2D slice, the current shape calculation is correct, and the tuple unpacking syntax is equivalent and cleaner:

-                indices = np.zeros((img_.shape[0],) + (img_.ndim - 1,) + img_.shape[1:], dtype=np.int32)
+                indices = np.zeros((img_.shape[0], img_.ndim - 1, *img_.shape[1:]), dtype=np.int32)

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

📥 Commits

Reviewing files that changed from the base of the PR and between 15fd428 and 258a472.

📒 Files selected for processing (1)

• monai/transforms/utils.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

⚙️ CodeRabbit configuration file

Review the Python code for quality and correctness. Ensure variable names adhere to PEP8 style guides, are sensible and informative in regards to their function, though permitting simple names for loop and comprehension variables. Ensure routine names are meaningful in regards to their function and use verbs, adjectives, and nouns in a semantically appropriate way. 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. Examine code for logical error or inconsistencies, and suggest what may be changed to addressed these. Suggest any enhancements for code improving efficiency, maintainability, comprehensibility, and correctness. Ensure new or modified definitions will be covered by existing or new unit tests.

Files:

• monai/transforms/utils.py

🪛 Ruff (0.14.8)

monai/transforms/utils.py

2501-2501: Consider (img.shape[0], img.dim() - 1, *img.shape[1:]) instead of concatenation

Replace with (img.shape[0], img.dim() - 1, *img.shape[1:])

(RUF005)

---

2535-2535: Consider (img_.shape[0], img_.ndim - 1, *img_.shape[1:]) instead of concatenation

Replace with (img_.shape[0], img_.ndim - 1, *img_.shape[1:])

(RUF005)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)

• GitHub Check: min-dep-py3 (3.12)
• GitHub Check: min-dep-py3 (3.11)
• GitHub Check: min-dep-pytorch (2.7.1)
• GitHub Check: min-dep-py3 (3.9)
• GitHub Check: min-dep-pytorch (2.5.1)
• GitHub Check: min-dep-pytorch (2.8.0)
• GitHub Check: min-dep-pytorch (2.6.0)
• GitHub Check: quick-py3 (macOS-latest)
• GitHub Check: min-dep-os (ubuntu-latest)
• GitHub Check: min-dep-py3 (3.10)
• GitHub Check: min-dep-os (windows-latest)
• GitHub Check: min-dep-os (macOS-latest)
• GitHub Check: flake8-py3 (mypy)
• GitHub Check: flake8-py3 (pytype)
• GitHub Check: quick-py3 (ubuntu-latest)
• GitHub Check: quick-py3 (windows-latest)
• GitHub Check: flake8-py3 (codeformat)
• GitHub Check: packaging
• GitHub Check: build-docs