Repository: keras-team/keras-cv

Issue:
RandomCutMix, RandomMixUp augmentation
These two well-known and very effective augmentation methods are widely used among ML practitioners. In keras, we can find some general augmentation . And these two augmentation methods can be included as advanced augmentation layers.

CutMix - [Paper](https://arxiv.org/abs/1905.04899) - Cited by ~ 865
MixUp - [Paper](https://arxiv.org/abs/1710.09412) - Cited by ~ 2675

Let's add these to 
`keras_cv.layers.RandomCutMix` and `keras_cv.layers.RandomMixUp`.

Discussion:
bhack: Just a preliminar note.
As we are using an avg lenght chain of TF ops how we are going to benchmarking these against performance/regression and the related  "transform/compiler stack" performance on this ops assembly (sub)graph?
LukeWood: > Just a preliminar note. As we are using an avg lenght chain of TF ops how we are going to benchmarking these against performance/regression and the related "transform/compiler stack" performance on this ops assembly (sub)graph?

Performance benchmarking is something I'm still working through, I'll be doing a thorough pass soon for my first implementation of COCO metrics.
bhack: Do you plan to [dump/inspect the compiled programs](https://www.tensorflow.org/xla#inspect_compiled_programs) and interact with the compiler team if some composition isn't handled efficently? Do you have reference performances to target? 
E.g. see more at https://discuss.tensorflow.org/t/deformable-convolution-and-other-custom-ops/1951


Or do you think to care only about some [e2e models use cases](https://github.com/tensorflow/benchmarks)?
LukeWood: Just discussed our plan here @bhack, we're going to test throughput of end to end models rather than individual components.  So at some point I'll put a classification model together that uses these KPLs
bhack: > Just discussed our plan here @bhack, we're going to test throughput of end to end models rather than individual components.  So at some point I'll put a classification model together that uses these KPLs

Sounds good for evaluating big O but we need to handle the model (size/computational overhead) selection with care. If not we could have some model bias about the expect impact of these components efficency.
LukeWood: Thanks for the review @fchollet!  I went ahead and fixed the issues. 
LukeWood: Appreciate the thorough review @qlzh727 and @fchollet.  Updated and I think the resulting class looks a lot better than it did before the review :).  Thanks!

Patch:
diff --git a/.gitignore b/.gitignore
index bc73860..5292914 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 keras_cv.egg-info/
+build/
 *.swp
 .idea
 *.pyc
diff --git a/examples/layers/preprocessing/cut_mix_demo.py b/examples/layers/preprocessing/cut_mix_demo.py
new file mode 100644
index 0000000..c6b4729
--- /dev/null
+++ b/examples/layers/preprocessing/cut_mix_demo.py
@@ -0,0 +1,47 @@
+"""cut_mix_demo.py shows how to use the CutMix preprocessing layer.
+
+Operates on the oxford_flowers102 dataset.  In this script the flowers 
+are loaded, then are passed through the preprocessing layers.  
+Finally, they are shown using matplotlib.
+"""
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from keras_cv.layers.preprocessing import cut_mix
+import matplotlib.pyplot as plt
+
+
+IMG_SIZE = (224, 224)
+BATCH_SIZE = 64
+
+
+def resize(image, label, num_classes=10):
+    image = tf.image.resize(image, IMG_SIZE)
+    label = tf.one_hot(label, num_classes)
+    return image, label
+
+
+def main():
+    data, ds_info = tfds.load("oxford_flowers102", with_info=True, as_supervised=True)
+    train_ds = data["train"]
+
+    num_classes = ds_info.features["label"].num_classes
+
+    train_ds = (
+        train_ds.map(lambda x, y: resize(x, y, num_classes=num_classes))
+        .shuffle(10 * BATCH_SIZE)
+        .batch(BATCH_SIZE)
+    )
+    cutmix = cut_mix.CutMix(1.0)
+    train_ds = train_ds.map(cutmix, num_parallel_calls=tf.data.AUTOTUNE)
+
+    for images, labels in train_ds.take(1):
+        plt.figure(figsize=(8, 8))
+        for i in range(9):
+            plt.subplot(3, 3, i + 1)
+            plt.imshow(images[i].numpy().astype("uint8"))
+            plt.axis("off")
+        plt.show()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/layers/preprocessing/mix_up_demo.py b/examples/layers/preprocessing/mix_up_demo.py
new file mode 100644
index 0000000..46b48f0
--- /dev/null
+++ b/examples/layers/preprocessing/mix_up_demo.py
@@ -0,0 +1,47 @@
+"""mix_up_demo.py shows how to use the MixUp preprocessing layer.
+
+Uses the oxford_flowers102 dataset.  In this script the flowers 
+are loaded, then are passed through the preprocessing layers.  
+Finally, they are shown using matplotlib.
+"""
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from keras_cv.layers.preprocessing import mix_up
+import matplotlib.pyplot as plt
+
+
+IMG_SIZE = (224, 224)
+BATCH_SIZE = 64
+
+
+def resize(image, label, num_classes=10):
+    image = tf.image.resize(image, IMG_SIZE)
+    label = tf.one_hot(label, num_classes)
+    return image, label
+
+
+def main():
+    data, ds_info = tfds.load("oxford_flowers102", with_info=True, as_supervised=True)
+    train_ds = data["train"]
+
+    num_classes = ds_info.features["label"].num_classes
+
+    train_ds = (
+        train_ds.map(lambda x, y: resize(x, y, num_classes=num_classes))
+        .shuffle(10 * BATCH_SIZE)
+        .batch(BATCH_SIZE)
+    )
+    mixup = mix_up.MixUp(1.0, alpha=0.8)
+    train_ds = train_ds.map(mixup, num_parallel_calls=tf.data.AUTOTUNE)
+
+    for images, labels in train_ds.take(1):
+        plt.figure(figsize=(8, 8))
+        for i in range(9):
+            plt.subplot(3, 3, i + 1)
+            plt.imshow(images[i].numpy().astype("uint8"))
+            plt.axis("off")
+        plt.show()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/keras_cv/__init__.py b/keras_cv/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/keras_cv/layers/__init__.py b/keras_cv/layers/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/keras_cv/layers/preprocessing/__init__.py b/keras_cv/layers/preprocessing/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/keras_cv/layers/preprocessing/cut_mix.py b/keras_cv/layers/preprocessing/cut_mix.py
new file mode 100644
index 0000000..c779bba
--- /dev/null
+++ b/keras_cv/layers/preprocessing/cut_mix.py
@@ -0,0 +1,174 @@
+import tensorflow as tf
+import tensorflow.keras as keras
+import tensorflow.keras.layers as layers
+from tensorflow.python.platform import tf_logging as logging
+
+
+class CutMix(layers.Layer):
+    """CutMix implements the CutMix data augmentation technique.
+
+    Args:
+        rate: Float between 0 and 1.  The fraction of samples to augment.
+        alpha: Float between 0 and 1.  Inverse scale parameter for the gamma distribution.
+            This controls the shape of the distribution from which the smoothing values are 
+            sampled.  Defaults 1.0, which is a recommended value when training an imagenet1k
+            classification model.
+        label_smoothing: Float in [0, 1]. When > 0, label values are smoothed, meaning the 
+            confidence on label values are relaxed. e.g. label_smoothing=0.2 means that we 
+            will use a value of 0.1 for label 0 and 0.9 for label 1.  Defaults 0.0.
+    References:
+       [CutMix paper]( https://arxiv.org/abs/1905.04899).
+
+    Sample usage:
+    ```python
+    (images, labels), _ = tf.keras.datasets.cifar10.load_data()
+    cutmix = keras_cv.layers.preprocessing.cut_mix.CutMix(10)
+    augmented_images, updated_labels = cutmix(images, labels)
+    ```
+    """
+
+    def __init__(
+        self, rate, label_smoothing=0.0, alpha=1.0, seed=None, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+        self.rate = rate
+        self.label_smoothing = label_smoothing
+        self.seed = seed
+
+    @staticmethod
+    def _sample_from_beta(alpha, beta, shape):
+        sample_alpha = tf.random.gamma(shape, 1.0, beta=alpha)
+        sample_beta = tf.random.gamma(shape, 1.0, beta=beta)
+        return sample_alpha / (sample_alpha + sample_beta)
+
+    def call(self, images, labels):
+        """call method for the CutMix layer.
+
+        Args:
+            images: Tensor representing images of shape [batch_size, width, height, channels], with dtype tf.float32.
+            labels: One hot encoded tensor of labels for the images, with dtype tf.float32.
+        Returns:
+            images: augmented images, same shape as input.
+            labels: updated labels with both label smoothing and the cutmix updates applied.
+        """
+
+        if tf.shape(images)[0] == 1:
+            logging.warning(
+                "CutMix received a single image to `call`.  The layer relies on combining multiple examples, "
+                "and as such will not behave as expected.  Please call the layer with 2 or more samples."
+            )
+
+        augment_cond = tf.less(
+            tf.random.uniform(shape=[], minval=0.0, maxval=1.0), self.rate
+        )
+        # pylint: disable=g-long-lambda
+        cutmix_augment = lambda: self._update_labels(*self._cutmix(images, labels))
+        no_augment = lambda: (images, self._smooth_labels(labels))
+        return tf.cond(augment_cond, cutmix_augment, no_augment)
+
+    def _cutmix(self, images, labels):
+        """Apply cutmix."""
+        input_shape = tf.shape(images)
+        batch_size, image_height, image_width = (
+            input_shape[0],
+            input_shape[1],
+            input_shape[2],
+        )
+
+        permutation_order = tf.random.shuffle(tf.range(0, batch_size), seed=self.seed)
+        lambda_sample = CutMix._sample_from_beta(self.alpha, self.alpha, (batch_size,))
+
+        ratio = tf.math.sqrt(1 - lambda_sample)
+
+        cut_height = tf.cast(
+            ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32
+        )
+        cut_width = tf.cast(
+            ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32
+        )
+
+        random_center_height = tf.random.uniform(
+            shape=[batch_size], minval=0, maxval=image_height, dtype=tf.int32
+        )
+        random_center_width = tf.random.uniform(
+            shape=[batch_size], minval=0, maxval=image_width, dtype=tf.int32
+        )
+
+        bbox_area = cut_height * cut_width
+        lambda_sample = 1.0 - bbox_area / (image_height * image_width)
+        lambda_sample = tf.cast(lambda_sample, dtype=tf.float32)
+
+        images = tf.map_fn(
+            lambda x: _fill_rectangle(*x),
+            (
+                images,
+                random_center_width,
+                random_center_height,
+                cut_width // 2,
+                cut_height // 2,
+                tf.gather(images, permutation_order),
+            ),
+            fn_output_signature=tf.TensorSpec.from_tensor(images[0]),
+        )
+
+        return images, labels, lambda_sample, permutation_order
+
+    def _update_labels(self, images, labels, lambda_sample, permutation_order):
+        labels_smoothed = self._smooth_labels(labels)
+        cutout_labels = tf.gather(labels, permutation_order)
+
+        lambda_sample = tf.reshape(lambda_sample, [-1, 1])
+        labels = lambda_sample * labels_smoothed + (1.0 - lambda_sample) * cutout_labels
+        return images, labels
+
+    def _smooth_labels(self, labels):
+        label_smoothing = self.label_smoothing or 0.0
+        off_value = label_smoothing / tf.cast(tf.shape(labels)[1], tf.float32)
+        on_value = 1.0 - label_smoothing + off_value
+        return labels * on_value + (1 - labels) * off_value
+
+
+def _fill_rectangle(
+    image, center_width, center_height, half_width, half_height, replace=None
+):
+    """Fill a rectangle in a given image using the value provided in replace.
+
+    Args:
+        image: the starting image to fill the rectangle on.
+        center_width: the X center of the rectangle to fill
+        center_height: the Y center of the rectangle to fill
+        half_width: 1/2 the width of the resulting rectangle
+        half_height: 1/2 the height of the resulting rectangle
+        replace: The value to fill the rectangle with.  Accepts a Tensor, 
+            Constant, or None.
+    Returns:
+        image: the modified image with the chosen rectangle filled.
+     """
+    image_shape = tf.shape(image)
+    image_height = image_shape[0]
+    image_width = image_shape[1]
+
+    lower_pad = tf.maximum(0, center_height - half_height)
+    upper_pad = tf.maximum(0, image_height - center_height - half_height)
+    left_pad = tf.maximum(0, center_width - half_width)
+    right_pad = tf.maximum(0, image_width - center_width - half_width)
+
+    cutout_shape = [
+        image_height - (lower_pad + upper_pad),
+        image_width - (left_pad + right_pad),
+    ]
+    padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
+    mask = tf.pad(
+        tf.zeros(cutout_shape, dtype=image.dtype), padding_dims, constant_values=1
+    )
+    mask = tf.expand_dims(mask, -1)
+
+    if replace is None:
+        fill = tf.random.normal(image_shape, dtype=image.dtype)
+    elif isinstance(replace, tf.Tensor):
+        fill = replace
+    else:
+        fill = tf.ones_like(image, dtype=image.dtype) * replace
+    image = tf.where(tf.equal(mask, 0), fill, image)
+    return image
diff --git a/keras_cv/layers/preprocessing/mix_up.py b/keras_cv/layers/preprocessing/mix_up.py
new file mode 100644
index 0000000..8f8d266
--- /dev/null
+++ b/keras_cv/layers/preprocessing/mix_up.py
@@ -0,0 +1,97 @@
+import tensorflow as tf
+import tensorflow.keras as keras
+import tensorflow.keras.layers as layers
+from tensorflow.python.platform import tf_logging as logging
+
+
+class MixUp(layers.Layer):
+    """MixUp implements the MixUp data augmentation technique.
+
+    Args:
+        rate: Float between 0 and 1.  The fraction of samples to augment.
+        alpha: Float between 0 and 1.  Inverse scale parameter for the gamma distribution.
+            This controls the shape of the distribution from which the smoothing values are 
+            sampled.  Defaults 0.2, which is a recommended value when training an imagenet1k
+            classification model.
+        label_smoothing: Float in [0, 1]. When > 0, label values are smoothed, meaning the 
+            confidence on label values are relaxed. e.g. label_smoothing=0.2 means that we 
+            will use a value of 0.1 for label 0 and 0.9 for label 1.  Defaults 0.0.
+    References:
+        [MixUp paper](https://arxiv.org/abs/1710.09412).
+
+    Sample usage:
+    ```python
+    (images, labels), _ = tf.keras.datasets.cifar10.load_data()
+    mixup = keras_cv.layers.preprocessing.mix_up.MixUp(10)
+    augmented_images, updated_labels = mixup(images, labels)
+    ```
+    """
+
+    def __init__(
+        self, rate, label_smoothing=0.0, alpha=0.2, seed=None, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.alpha = alpha
+        self.rate = rate
+        self.label_smoothing = label_smoothing
+        self.seed = seed
+
+    @staticmethod
+    def _sample_from_beta(alpha, beta, shape):
+        sample_alpha = tf.random.gamma(shape, 1.0, beta=alpha)
+        sample_beta = tf.random.gamma(shape, 1.0, beta=beta)
+        return sample_alpha / (sample_alpha + sample_beta)
+
+    def call(self, images, labels):
+        """call method for the MixUp layer.
+
+        Args:
+            images: Tensor representing images of shape [batch_size, width, height, channels], with dtype tf.float32.
+            labels: One hot encoded tensor of labels for the images, with dtype tf.float32.
+        Returns:
+            images: augmented images, same shape as input.
+            labels: updated labels with both label smoothing and the cutmix updates applied.
+        """
+
+        if tf.shape(images)[0] == 1:
+            logging.warning(
+                "MixUp received a single image to `call`.  The layer relies on combining multiple examples, "
+                "and as such will not behave as expected.  Please call the layer with 2 or more samples."
+            )
+
+        augment_cond = tf.less(
+            tf.random.uniform(shape=[], minval=0.0, maxval=1.0), self.rate
+        )
+        # pylint: disable=g-long-lambda
+        mixup_augment = lambda: self._update_labels(*self._mixup(images, labels))
+        no_augment = lambda: (images, self._smooth_labels(labels))
+        return tf.cond(augment_cond, mixup_augment, no_augment)
+
+    def _mixup(self, images, labels):
+        batch_size = tf.shape(images)[0]
+        permutation_order = tf.random.shuffle(tf.range(0, batch_size), seed=self.seed)
+
+        lambda_sample = MixUp._sample_from_beta(self.alpha, self.alpha, (batch_size,))
+        lambda_sample = tf.reshape(lambda_sample, [-1, 1, 1, 1])
+
+        mixup_images = tf.gather(images, permutation_order)
+        images = lambda_sample * images + (1.0 - lambda_sample) * mixup_images
+
+        return images, labels, tf.squeeze(lambda_sample), permutation_order
+
+    def _update_labels(self, images, labels, lambda_sample, permutation_order):
+        labels_smoothed = self._smooth_labels(labels)
+        labels_for_mixup = tf.gather(labels, permutation_order)
+
+        lambda_sample = tf.reshape(lambda_sample, [-1, 1])
+        labels = (
+            lambda_sample * labels_smoothed + (1.0 - lambda_sample) * labels_for_mixup
+        )
+
+        return images, labels
+
+    def _smooth_labels(self, labels):
+        label_smoothing = self.label_smoothing or 0.0
+        off_value = label_smoothing / tf.cast(tf.shape(labels)[1], tf.float32)
+        on_value = 1.0 - label_smoothing + off_value
+        return on_value * labels + (1 - labels) * off_value
diff --git a/setup.py b/setup.py
index 4312819..9facabb 100644
--- a/setup.py
+++ b/setup.py
@@ -25,7 +25,10 @@ setup(
     author_email="keras-cv@google.com",
     license="Apache License 2.0",
     install_requires=["packaging", "tensorflow"],
-    extras_require={"tests": ["flake8", "isort", "black",],},
+    extras_require={
+        "tests": ["flake8", "isort", "black",],
+        "examples": ["tensorflow_datasets", "matplotlib"],
+    },
     classifiers=[
         "Programming Language :: Python",
         "Programming Language :: Python :: 3.7",


Test patch:
diff --git a/keras_cv/layers/preprocessing/cut_mix_test.py b/keras_cv/layers/preprocessing/cut_mix_test.py
new file mode 100644
index 0000000..a917350
--- /dev/null
+++ b/keras_cv/layers/preprocessing/cut_mix_test.py
@@ -0,0 +1,91 @@
+import tensorflow as tf
+from keras_cv.layers.preprocessing.cut_mix import CutMix
+
+
+NUM_CLASSES = 10
+
+
+class CutMixTest(tf.test.TestCase):
+    def test_return_shapes(self):
+        xs = tf.ones((2, 512, 512, 3))
+        # randomly sample labels
+        ys = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
+        ys = tf.squeeze(ys)
+        ys = tf.one_hot(ys, NUM_CLASSES)
+
+        layer = CutMix(1.0, seed=1)
+        xs, ys = layer(xs, ys)
+
+        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        # one hot smoothed labels
+        self.assertEqual(ys.shape, [2, 10])
+        self.assertEqual(len(ys != 0.0), 2)
+
+    def test_label_smoothing(self):
+        xs = tf.ones((2, 512, 512, 3))
+        # randomly sample labels
+        ys = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
+        ys = tf.squeeze(ys)
+        ys = tf.one_hot(ys, NUM_CLASSES)
+
+        layer = CutMix(1.0, label_smoothing=0.2, seed=1)
+        xs, ys = layer(xs, ys)
+        self.assertNotAllClose(ys, 0.0)
+
+    def test_cut_mix_call_results(self):
+        xs = tf.cast(
+            tf.stack([2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))], axis=0,), tf.float32,
+        )
+        ys = tf.one_hot(tf.constant([0, 1]), 2)
+
+        layer = CutMix(1.0, label_smoothing=0.0, seed=1)
+        xs, ys = layer(xs, ys)
+
+        # At least some pixels should be replaced in the CutMix operation
+        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        # No labels should still be close to their original values
+        self.assertNotAllClose(ys, 1.0)
+        self.assertNotAllClose(ys, 0.0)
+
+    def test_cut_mix_call_results_one_channel(self):
+        xs = tf.cast(
+            tf.stack([2 * tf.ones((4, 4, 1)), tf.ones((4, 4, 1))], axis=0,), tf.float32,
+        )
+        ys = tf.one_hot(tf.constant([0, 1]), 2)
+
+        layer = CutMix(1.0, label_smoothing=0.0, seed=1)
+        xs, ys = layer(xs, ys)
+
+        # At least some pixels should be replaced in the CutMix operation
+        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        # No labels should still be close to their original values
+        self.assertNotAllClose(ys, 1.0)
+        self.assertNotAllClose(ys, 0.0)
+
+    def test_in_tf_function(self):
+        xs = tf.cast(
+            tf.stack([2 * tf.ones((100, 100, 1)), tf.ones((100, 100, 1))], axis=0), tf.float32,
+        )
+        ys = tf.one_hot(tf.constant([0, 1]), 2)
+
+        layer = CutMix(1.0, label_smoothing=0.0, seed=1)
+
+        @tf.function
+        def augment(x, y):
+            return layer(x, y)
+        xs, ys = augment(xs, ys)
+
+        # At least some pixels should be replaced in the CutMix operation
+        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        # No labels should still be close to their original values
+        self.assertNotAllClose(ys, 1.0)
+        self.assertNotAllClose(ys, 0.0)
diff --git a/keras_cv/layers/preprocessing/mix_up_test.py b/keras_cv/layers/preprocessing/mix_up_test.py
new file mode 100644
index 0000000..b648ce8
--- /dev/null
+++ b/keras_cv/layers/preprocessing/mix_up_test.py
@@ -0,0 +1,71 @@
+import tensorflow as tf
+from keras_cv.layers.preprocessing.mix_up import MixUp
+
+
+NUM_CLASSES = 10
+
+
+class MixUpTest(tf.test.TestCase):
+    def test_return_shapes(self):
+        xs = tf.ones((2, 512, 512, 3))
+        # randomly sample labels
+        ys = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
+        ys = tf.squeeze(ys)
+        ys = tf.one_hot(ys, NUM_CLASSES)
+
+        layer = MixUp(1.0)
+        xs, ys = layer(xs, ys)
+
+        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        # one hot smoothed labels
+        self.assertEqual(ys.shape, [2, 10])
+        self.assertEqual(len(ys != 0.0), 2)
+
+    def test_label_smoothing(self):
+        xs = tf.ones((2, 512, 512, 3))
+        # randomly sample labels
+        ys = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
+        ys = tf.squeeze(ys)
+        ys = tf.one_hot(ys, NUM_CLASSES)
+
+        layer = MixUp(1.0, label_smoothing=0.2)
+        xs, ys = layer(xs, ys)
+        self.assertNotAllClose(ys, 0.0)
+
+    def test_cut_mix_call_results(self):
+        xs = tf.cast(
+            tf.stack([2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))], axis=0,), tf.float32,
+        )
+        ys = tf.one_hot(tf.constant([0, 1]), 2)
+
+        layer = MixUp(1.0, label_smoothing=0.0)
+        xs, ys = layer(xs, ys)
+
+        # None of the individual values should still be close to 1 or 0
+        self.assertNotAllClose(xs, 1.0)
+        self.assertNotAllClose(xs, 2.0)
+
+        # No labels should still be close to their originals
+        self.assertNotAllClose(ys, 1.0)
+        self.assertNotAllClose(ys, 0.0)
+
+    def test_in_tf_function(self):
+        xs = tf.cast(
+            tf.stack([2 * tf.ones((4, 4, 3)), tf.ones((4, 4, 3))], axis=0,), tf.float32,
+        )
+        ys = tf.one_hot(tf.constant([0, 1]), 2)
+
+        layer = MixUp(1.0, label_smoothing=0.0)
+
+        @tf.function
+        def augment(x, y):
+            return layer(x, y)
+        xs, ys = augment(xs, ys)
+
+        # None of the individual values should still be close to 1 or 0
+        self.assertNotAllClose(xs, 1.0)
+        self.assertNotAllClose(xs, 2.0)
+
+        # No labels should still be close to their originals
+        self.assertNotAllClose(ys, 1.0)
+        self.assertNotAllClose(ys, 0.0)