Add channel_axis to random_flip_up_down

dm_pix/_src/augment.py

@@ -803,6 +803,7 @@ def random_flip_up_down(
     image: chex.Array,
     *,
     probability: chex.Numeric = 0.5,
+    channel_axis: int = -1,
 ) -> chex.Array:
   """Applies `flip_up_down` with a given probability.
 
@@ -812,14 +813,20 @@ def random_flip_up_down(
       ...HWC or ...CHW.
     probability: the probability of applying flip_up_down transform. Must be a
       value in [0, 1].
+    channel_axis: the index of the channel axis.
 
   Returns:
     An up-down flipped image if condition is met, otherwise original image.
   """
   # DO NOT REMOVE - Logging usage.
 
   should_transform = jax.random.bernoulli(key=key, p=probability)
-  return jax.lax.cond(should_transform, flip_up_down, lambda x: x, image)
+  return jax.lax.cond(
+      should_transform,
+      lambda x: flip_up_down(x, channel_axis=channel_axis),
+      lambda x: x,
+      image,
+  )
 
 
 def random_brightness(

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_up_down_channel_axis(self):
+    image = jnp.arange(2 * 3 * 4, dtype=jnp.float32).reshape((3, 2, 4))
+    result = augment.random_flip_up_down(
+        jax.random.PRNGKey(0), image, probability=1., channel_axis=0)
+    np.testing.assert_array_equal(result, jnp.flip(image, axis=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(