[krakken190]

Fixes #141.

Problem

rotate and affine_transform say they support HWC or CHW images, but they always treat the first two array axes as height and width. For CHW inputs this rotates or shifts across the channel axis and silently returns a corrupted image.

Fix

• Add channel_axis: int = -1 to affine_transform. When it is not -1, the channel axis is moved to the end before the transform and moved back afterwards. This is the same pattern elastic_deformation uses.
• The matrix and offset are always interpreted in HWC order, independent of channel_axis, so one transform can be applied to either layout. The docstring now says this.
• Add channel_axis: int = -1 to rotate. It computes the rotation centre from the HWC-ordered shape and forwards channel_axis to affine_transform.
• The default -1 keeps the existing behaviour, so current callers are unaffected.

Tests

Added to TestCustom:

• test_affine_transform_channel_axis: for nearest/0, nearest/1, mirror/1 and constant/1, with channel_axis 0 and -3, the CHW result equals the HWC result moved to channels-first.
• test_rotate_channel_axis: the same check for two angles, run under jax.jit.

Validation

• pytest dm_pix/_src/augment_test.py: 886 passed
• pytest dm_pix/_src/api_test.py: 80 passed
• Manual check on a real image: with channel_axis=0, CHW output matches HWC exactly (max diff 0.0) for both functions. Without it, the difference was 0.99997 before and after the change.

Tested on CPU with Python 3.12.2, JAX/jaxlib 0.12.0, NumPy 2.5.4, SciPy 1.18.1 and TensorFlow 2.21.0. GPU/TPU not tested.

[claudiofantacci]

Thanks!

[krakken190]

Hi @claudiofantacci , I've resolved the merge conflict with master. Ready for review/merge whenever you get a chance!

[krakken190]

@claudiofantacci Sorry, my merge-conflict resolution in the GitHub UI dropped the @parameterized.product( decorator above test_gaussian_blur_channel_axis, which caused the IndentationError in CI on 72a3195. That's fixed in e76e336. The branch now differs from master only by the original change you approved, and augment_test.py, color_conversion_test.py and api_test.py pass locally (1192 passed). Could you approve the workflows again when you have a moment? Thanks!