def METHOD_NAME(self):
    xs = tf.ones((2, 512, 512, 3))
    # randomly sample labels
    ys_labels = tf.random.categorical(tf.math.log([[0.5, 0.5]]), 2)
    ys_labels = tf.squeeze(ys_labels)
    ys_labels = tf.one_hot(ys_labels, num_classes)
    # randomly sample bounding boxes
    ys_bounding_boxes = {
        "boxes": tf.random.uniform((2, 3, 4), 0, 1),
        "classes": tf.random.uniform((2, 3), 0, 1),
    }
    # randomly sample segmentation mask
    ys_segmentation_masks = tf.cast(
        tf.stack(
            [2 * tf.ones((512, 512)), tf.ones((512, 512))],
            axis=0,
        ),
        tf.uint8,
    )
    ys_segmentation_masks = tf.one_hot(ys_segmentation_masks, 3)
    layer = MixUp()
    # mixup on labels
    outputs = layer(
        {
            "images": xs,
            "labels": ys_labels,
            "bounding_boxes": ys_bounding_boxes,
            "segmentation_masks": ys_segmentation_masks,
        }
    )
    xs, ys_labels, ys_bounding_boxes, ys_segmentation_masks = (
        outputs["images"],
        outputs["labels"],
        outputs["bounding_boxes"],
        outputs["segmentation_masks"],
    )
    self.assertEqual(xs.shape, [2, 512, 512, 3])
    self.assertEqual(ys_labels.shape, [2, 10])
    self.assertEqual(ys_bounding_boxes["boxes"].shape, [2, 6, 4])
    self.assertEqual(ys_bounding_boxes["classes"].shape, [2, 6])
    self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])