From 37735ae4cf4d43c0414da36d8479675a096a38a2 Mon Sep 17 00:00:00 2001 From: v1shay Date: Wed, 12 Aug 2026 11:20:40 -0700 Subject: [PATCH 1/2] Add channel axis to random horizontal flip --- dm_pix/_src/augment.py | 9 ++++++++- dm_pix/_src/augment_test.py | 6 ++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index adad96f..a419e52 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -772,6 +772,7 @@ def random_flip_left_right( image: chex.Array, *, probability: chex.Numeric = 0.5, + channel_axis: int = -1, ) -> chex.Array: """Applies `flip_left_right` with a given probability. @@ -781,6 +782,7 @@ def random_flip_left_right( ...HWC or ...CHW. probability: the probability of applying flip_left_right transform. Must be a value in [0, 1]. + channel_axis: the index of the channel axis. Returns: A left-right flipped image if condition is met, otherwise original image. @@ -788,7 +790,12 @@ def random_flip_left_right( # DO NOT REMOVE - Logging usage. should_transform = jax.random.bernoulli(key=key, p=probability) - return jax.lax.cond(should_transform, flip_left_right, lambda x: x, image) + return jax.lax.cond( + should_transform, + lambda x: flip_left_right(x, channel_axis=channel_axis), + lambda x: x, + image, + ) def random_flip_up_down( diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 8532e94..e83c7f4 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -188,6 +188,12 @@ def test_flip(self, images_list): reference_fn=None, probability=(0., 1.)) + def test_random_flip_left_right_channel_axis(self): + image = jnp.arange(2 * 3 * 4, dtype=jnp.float32).reshape((3, 2, 4)) + result = augment.random_flip_left_right( + jax.random.PRNGKey(0), image, probability=1., channel_axis=0) + np.testing.assert_array_equal(result, jnp.flip(image, axis=2)) + # Due to a bug in scipy we cannot test all available modes, refer to these # issues for more information: https://github.com/jax-ml/jax/issues/11097, # https://github.com/jax-ml/jax/issues/11097 From faf70bc7e8a1d4adbd4f052d5ffef5e66c5daefb Mon Sep 17 00:00:00 2001 From: v1shay Date: Wed, 12 Aug 2026 13:34:46 -0700 Subject: [PATCH 2/2] Re-run CLA check