def forward(self, x: "tf.Tensor", y: "tf.Tensor" | None = None) -> tuple["tf.Tensor", "tf.Tensor" | None]:
        """
        Apply CutMix data augmentation to sample `x`.

        :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
        :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
                  or class indices of shape `(nb_samples,)`.
        :return: Data augmented sample. The returned labels will be probability vectors of shape
                 `(nb_samples, nb_classes)`.
        """
        import tensorflow as tf

        if y is None:
            raise ValueError("Labels `y` cannot be None.")

        # convert labels to one-hot encoding
        if len(y.shape) == 2:
            y_one_hot = y
        elif len(y.shape) == 1:
            y_one_hot = tf.one_hot(y, self.num_classes, on_value=1.0, off_value=0.0)
        else:
            raise ValueError(
                "Shape of labels not recognised. "
                "Please provide labels in shape (nb_samples,) or (nb_samples, nb_classes)"
            )

        x_ndim = len(x.shape)

        # NCHW/NHWC/NCFHW --> NFHWC
        if x_ndim == 4:
            if self.channels_first:
                # NCHW --> NHWC --> NFHWC
                x_nhwc = tf.transpose(x, (0, 2, 3, 1))
                x_nfhwc = tf.expand_dims(x_nhwc, axis=1)
            else:
                # NHWC --> NFHWC
                x_nfhwc = tf.expand_dims(x, axis=1)
        elif x_ndim == 5:
            if self.channels_first:
                # NCFHW --> NFHWC
                x_nfhwc = tf.transpose(x, (0, 2, 3, 4, 1))
            else:
                # NFHWC
                x_nfhwc = x
        else:
            raise ValueError("Unrecognized input dimension. CutMix can only be applied to image and video data.")

        n, _, height, width, _ = x_nfhwc.shape
        x_aug = tf.Variable(x_nfhwc, trainable=False)
        y_aug = tf.Variable(y_one_hot, trainable=False)

        # sample the combination ratio from the Beta distribution
        lmb = np.random.beta(self.alpha, self.alpha)
        cut_ratio = np.sqrt(1.0 - lmb)
        cut_height = int(height * cut_ratio)
        cut_width = int(width * cut_ratio)

        # randomly choose indices for samples to mix
        indices = tf.random.shuffle(tf.range(n))

        # generate a random bounding box per image
        for idx1, idx2 in enumerate(tqdm(indices, desc="CutMix", disable=not self.verbose)):
            prob = np.random.rand()
            if prob < self.probability:
                # uniform sampling
                center_y = tf.random.uniform(shape=[], maxval=height, dtype=tf.int32)
                center_x = tf.random.uniform(shape=[], maxval=width, dtype=tf.int32)
                bby1 = tf.clip_by_value(center_y - cut_height // 2, 0, height)
                bbx1 = tf.clip_by_value(center_x - cut_width // 2, 0, width)
                bby2 = tf.clip_by_value(center_y + cut_height // 2, 0, height)
                bbx2 = tf.clip_by_value(center_x + cut_width // 2, 0, width)

                # insert image bounding box
                x_aug[idx1, :, bbx1:bbx2, bby1:bby2, :].assign(x_nfhwc[idx2, :, bbx1:bbx2, bby1:bby2, :])
                # mix labels
                y_aug[idx1].assign(lmb * y_aug[idx1] + (1.0 - lmb) * y_one_hot[idx2])

        x_nfhwc = tf.convert_to_tensor(x_aug)
        y_aug = tf.convert_to_tensor(y_aug)

        # NCHW/NHWC/NCFHW <-- NFHWC
        if x_ndim == 4:
            if self.channels_first:
                # NHWC <-- NCHW <-- NFHWC
                x_nhwc = tf.squeeze(x_nfhwc, axis=1)
                x_aug = tf.transpose(x_nhwc, (0, 3, 1, 2))
            else:
                # NHWC <-- NFHWC
                x_aug = tf.squeeze(x_nfhwc, axis=1)
        elif x_ndim == 5:
            if self.channels_first:
                # NCFHW <-- NFHWC
                x_aug = tf.transpose(x_nfhwc, (0, 4, 1, 2, 3))
            else:
                # NFHWC
                x_aug = x_nfhwc

        return x_aug, y_aug

        Apply CutMix data augmentation to sample `x`.

        :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
        :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
                  or class indices of shape `(nb_samples,)`.
        :return: Data augmented sample. The returned labels will be probability vectors of shape
                 `(nb_samples, nb_classes)`.
        
art/defences/preprocessor/cutmix/cutmix_tensorflow.py