Fix bug when upranking passthrough inputs to RandAugment (#2194)

- RandAugment sometimes will choose a "no augmentation" option and
  passthrough inputs unaltered.
- Preprocessing normalization routines were not making copies of inputs
  and sometimes mutating layer input directly (mutating the input
  dict to cast dtypes and uprank tensors).
- RandAugment under the passthrough option would return these inputs
  directly.

The net effect was sometimes attempting to uprank during a passthrough
call, breaking tf.map_fn

Author: mattdangerw

keras_cv/layers/preprocessing/base_image_augmentation_layer.py

@@ -571,6 +571,8 @@ def _ensure_inputs_are_compute_dtype(self, inputs):
                 inputs,
                 self.compute_dtype,
             )
+        # Copy the input dict before we mutate it.
+        inputs = dict(inputs)
         inputs[IMAGES] = preprocessing.ensure_tensor(
             inputs[IMAGES],
             self.compute_dtype,

keras_cv/layers/preprocessing/rand_augment_test.py

@@ -12,17 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
-import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
 from keras_cv import layers
-from keras_cv.backend.config import keras_3
 from keras_cv.tests.test_case import TestCase
 
 
-@pytest.mark.skipif(keras_3(), reason="imcompatible with Keras 3")
 class RandAugmentTest(TestCase):
+    def test_zero_rate_pass_through(self):
+        rand_augment = layers.RandAugment(
+            value_range=(0, 255),
+            rate=0.0,
+        )
+        xs = np.ones((2, 512, 512, 3))
+        ys = rand_augment(xs)
+        self.assertAllClose(ys, xs)
+
     @parameterized.named_parameters(
         ("0", 0),
         ("20", 0.2),

keras_cv/layers/preprocessing/random_augmentation_pipeline.py

@@ -103,7 +103,7 @@ def _augment(self, inputs):
             )
             result = tf.cond(
                 skip_augment > self.rate,
-                lambda: inputs,
+                lambda: result,
                 lambda: self._random_choice(result),
             )
         return result

keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py

@@ -444,6 +444,9 @@ def _format_inputs(self, inputs):
             # single image input tensor
             metadata[IS_DICT] = False
             inputs = {IMAGES: inputs}
+        else:
+            # Copy the input dict before we mutate it.
+            inputs = dict(inputs)
 
         metadata[BATCHED] = inputs["images"].shape.rank == 4
         if inputs["images"].shape.rank == 3:
@@ -504,6 +507,8 @@ def _ensure_inputs_are_compute_dtype(self, inputs):
                 inputs,
                 self.compute_dtype,
             )
+        # Copy the input dict before we mutate it.
+        inputs = dict(inputs)
         inputs[IMAGES] = preprocessing.ensure_tensor(
             inputs[IMAGES],
             self.compute_dtype,