Overview
I contributed a fix to the PyTorch codebase addressing an edge case in torch.unravel_index, and the change was reviewed and merged into PyTorch's main branch.
torch.unravel_index converts flat indices into coordinates for a tensor of a given shape. While investigating its behavior around zero-sized dimensions, I found a case where invalid input could propagate into the underlying computation and result in division/modulo by zero rather than being rejected with a meaningful validation error.
Merged PR: https://github.com/pytorch/pytorch/pull/191092
The Problem
Consider a tensor shape containing a zero-sized dimension while providing non-empty indices.
A shape containing a zero dimension represents a tensor with zero total elements. Therefore, there cannot be a valid non-empty flat index into that shape.
However, this edge case could make it into the coordinate computation performed by torch.unravel_index. Because the implementation relies on arithmetic involving the dimensions of the shape, a zero dimension could eventually result in a division or modulo-by-zero operation.
Instead of clearly communicating that the supplied indices were invalid for the requested shape, the operation could therefore fail with an unexpected runtime error.
The Fix
I added explicit validation for this case before the coordinate calculation takes place.
The resulting behavior distinguishes between two important cases:
-
Zero-sized shape + non-empty indices: The input is invalid and now raises a clear
ValueError. - Zero-sized shape + empty indices: The operation remains valid and produces empty coordinate tensors.
Preserving the second case was important because an empty collection of indices does not attempt to reference an element that doesn't exist. It also maintains behavior consistent with NumPy for the corresponding edge case.
Reasoning Through the Edge Case
One of the most interesting parts of the contribution came during code review.
A question was raised about whether the validation should depend on all dimensions or whether the presence of any zero-sized dimension was sufficient to make non-empty indices invalid.
Working through the implementation made the answer clearer.
If any dimension of a tensor's shape is zero, then:
numel = shape[0] × shape[1] × ... × shape[n] = 0
The tensor therefore contains no elements.
Consequently, there is no valid flat index into that tensor, regardless of where the zero appears in its shape.
Tracing the implementation also required thinking through the coefficient calculations, broadcasting behavior, division, and modulo operations used to transform flat indices into multidimensional coordinates.
Testing
Along with the validation change, I added regression coverage for the edge case to ensure that invalid non-empty indices are rejected while the valid empty-index behavior continues to work correctly.
The change then went through PyTorch's upstream review and CI process before being approved and merged.
What I Learned
The code change itself is relatively small, but working in a production ML framework made the surrounding engineering particularly valuable.
I gained experience with:
- Navigating a large open-source codebase like PyTorch
- Understanding implementation details behind a familiar tensor API
- Reasoning carefully about tensor shapes and zero-sized dimensions
- Designing validation around API semantics rather than simply preventing a crash
- Writing regression tests for unusual edge cases
- Responding to technical feedback during upstream code review
- Working through the contribution and CI process of a major open-source project
One of my biggest takeaways was that robustness in foundational libraries often comes from handling seemingly tiny edge cases correctly. When an API is used as widely as PyTorch, even input validation deserves careful consideration around semantics, compatibility, and existing behavior.
I'm excited to continue contributing to open source and exploring more of the systems and infrastructure behind machine learning.
Links
PyTorch PR #191092:
https://github.com/pytorch/pytorch/pull/191092
Technologies: Python · PyTorch · NumPy · Git · GitHub · Open Source · Testing
Top comments (0)