random_flip_left_right does not support HWC

[Logon27]

https://dm-pix.readthedocs.io/en/latest/api.html#random-flip-left-right
Despite being documented to say so, there is no way to actually pass in the channel_axis to this function. So you cannot set the image format to be HWC. Please correct me if I am wrong.

[claudiofantacci]

Hi @Logon27, thanks for reaching out 😄

So the function here

dm_pix/dm_pix/_src/augment.py

Line 788 in ae47829

return jax.lax.cond(should_transform, flip_left_right, lambda x: x, image)

calls flip_left_right in its default mode which is this one

dm_pix/dm_pix/_src/augment.py

Lines 415 to 419 in ae47829

	def flip_left_right(
	image: chex.Array,
	*,
	channel_axis: int = -1,
	) -> chex.Array:

hence it is assumed that the image is of shape ...HWC.

There's indeed a missing piping, and that's true for all random functions actually.

Feel free to submit a PR with the changes and some test coverage if you want to!