def METHOD_NAME(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 = FourierMix()
    @tf.function
    def augment(x, y):
        return layer({"images": x, "labels": y})
    outputs = augment(xs, ys)
    xs, ys = outputs["images"], outputs["labels"]
    # 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)