def forward(self, x: "tf.Tensor", y: "tf.Tensor" | None = None) -> tuple["tf.Tensor", "tf.Tensor" | None]:
        """
        Apply local spatial smoothing to sample `x`.
        """
        import tensorflow as tf
        import tensorflow_addons as tfa

        x_ndim = x.ndim

        if x_ndim == 4:
            x_nhwc = x
        elif x_ndim == 5:
            # NFHWC --> NHWC
            nb_clips, clip_size, height, width, channels = x.shape
            x_nhwc = tf.reshape(x, (nb_clips * clip_size, height, width, channels))
        else:
            raise ValueError(
                "Unrecognized input dimension. Spatial smoothing can only be applied to image (NHWC) and video (NFHWC) "
                "data."
            )

        x_nhwc = tfa.image.median_filter2d(
            x_nhwc, filter_shape=[self.window_size, self.window_size], padding="REFLECT", constant_values=0, name=None
        )

        if x_ndim == 4:
            x = x_nhwc
        elif x_ndim == 5:
            # NFHWC <-- NHWC
            x = tf.reshape(x_nhwc, (nb_clips, clip_size, height, width, channels))

        if self.clip_values is not None:
            x = x.clip_by_value(min=self.clip_values[0], max=self.clip_values[1])

        return x, y

        Apply local spatial smoothing to sample `x`.
        
art/defences/preprocessor/spatial_smoothing_tensorflow.py