fix: respect channel_axis in pad_to_size for channels-first images#112
Merged
copybara-service[bot] merged 2 commits intoJun 2, 2026
Conversation
claudiofantacci
requested changes
Jun 1, 2026
Collaborator
claudiofantacci
left a comment
claudiofantacci
left a comment
There was a problem hiding this comment.
Thanks a lot! 🚀
| import numpy as np | ||
|
|
||
|
|
||
| class PadToSizeChannelsFirstTest(absltest.TestCase): |
Collaborator
There was a problem hiding this comment.
Please move this test under dm_pix/_src/augment_test.py.
| def test_pad_to_size_chw_shape(self): | ||
| """Channels-first (CHW) with channel_axis=0: output shape must be (C, H', W'). | ||
|
|
||
| Bug: pad_to_size hardcodes pad_width as ((top,bot), (left,right), (0,0)), |
Collaborator
There was a problem hiding this comment.
Remove this comment please.
Nas01010101
added a commit
to Nas01010101/dm_pix
that referenced
this pull request
Jun 1, 2026
Address review feedback on PR google-deepmind#112: - Move the channels-first pad_to_size regression tests out of the standalone pad_to_size_test.py and into the existing TestCustom class in augment_test.py, matching that file's conventions. - Drop the explanatory comments, consistent with the surrounding tests. Delete dm_pix/_src/pad_to_size_test.py.
Collaborator
claudiofantacci
left a comment
claudiofantacci
left a comment
There was a problem hiding this comment.
Thanks! It will take a little while to go through and I actually need to drop support from some old Python version.
I'll approve once this is done and then go from there.
In the meantime, thanks again!
Collaborator
|
Waiting #113 to go through. |
Contributor
Author
|
Sounds good, thanks!! |
Collaborator
|
Can you please rebase on top of main after #113 and re-upload? |
Contributor
Author
|
Of course ! |
pad_to_size hardcoded pad_width as ((top, bottom), (left, right), (0, 0)) which assumes channels-last (HWC) layout regardless of the channel_axis argument. For channels-first (CHW) images, this causes the padding to be applied along the channel dimension and leaves the width dimension unpadded, producing a wrong output shape (e.g. (5, 6, 4) instead of (3, 6, 6)). Fix: branch on _channels_last(image, channel_axis) and produce ((0, 0), (top, bottom), (left, right)) for CHW layout, consistent with how every other function in this module handles channel_axis. Add pad_to_size_test.py covering: - HWC baseline (channel_axis=-1): shape is preserved correctly. - CHW shape (channel_axis=0): must be (C, H', W') not (C+pad, H', W). - CHW values (channel_axis=0): original pixels centered, border is zero. - Batched CHW (channel_axis=1): must be (B, C, H', W').
Address review feedback on PR google-deepmind#112: - Move the channels-first pad_to_size regression tests out of the standalone pad_to_size_test.py and into the existing TestCustom class in augment_test.py, matching that file's conventions. - Drop the explanatory comments, consistent with the surrounding tests. Delete dm_pix/_src/pad_to_size_test.py.
Nas01010101
force-pushed
the
fix/pad-to-size-channels-first
branch
from
6617907 to
5ab629f
Compare
claudiofantacci
approved these changes
Jun 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
pad_to_size()ignoreschannel_axiswhen constructing the pad widths, so it pads the wrong axes for channels-first (CHW) images.The padding amounts (
top,bottom,left,right) are computed correctly via_get_dimension_values, which is channel-axis-aware. But thepad_widthtuple passed tojnp.padis hardcoded to channels-last order:For a CHW image (
channel_axis=0) this applies the height/width padding to the channel and height axes and leaves width unpadded. Example — padding a(3, 4, 4)CHW image to a(6, 6)target yields(5, 6, 4)instead of the expected(3, 6, 6).Every other spatial op in this module already branches on
_channels_last(image, channel_axis);pad_to_sizesimply missed it.Fix
Select the pad-width tuple based on layout:
The batch-dimension prepend and all other logic are layout-agnostic and unchanged. Channels-last behaviour is identical to before.
Tests
Adds
dm_pix/_src/pad_to_size_test.pycovering HWC (unchanged), CHW shape, CHW values, and batched CHW. The CHW tests fail before the fix and pass after.Broader regression (
interpolation_test.py+ the new test): 100 passed, 24 skipped (skips are pre-existing TF-gated tests).