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

• [x] 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.
• [x] Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• [x] 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[bot]]

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.}