src/transformers/models/mobilevit/modeling_tf_mobilevit.py/TFMobileViTLayer/__init__
class TFMobileViTLayer:    def __init__(
        self,
        config: MobileViTConfig,
        in_channels: int,
        out_channels: int,
        stride: int,
        hidden_size: int,
        num_stages: int,
        dilation: int = 1,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.patch_width = config.patch_size
        self.patch_height = config.patch_size

        if stride == 2:
            self.downsampling_layer = TFMobileViTInvertedResidual(
                config,
                in_channels=in_channels,
                out_channels=out_channels,
                stride=stride if dilation == 1 else 1,
                dilation=dilation // 2 if dilation > 1 else 1,
                name="downsampling_layer",
            )
            in_channels = out_channels
        else:
            self.downsampling_layer = None

        self.conv_kxk = TFMobileViTConvLayer(
            config, out_channels=in_channels, kernel_size=config.conv_kernel_size, name="conv_kxk"
        )

        self.conv_1x1 = TFMobileViTConvLayer(
            config,
            out_channels=hidden_size,
            kernel_size=1,
            use_normalization=False,
            use_activation=False,
            name="conv_1x1",
        )

        self.transformer = TFMobileViTTransformer(
            config, hidden_size=hidden_size, num_stages=num_stages, name="transformer"
        )

        self.layernorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")

        self.conv_projection = TFMobileViTConvLayer(
            config, out_channels=in_channels, kernel_size=1, name="conv_projection"
        )

        self.fusion = TFMobileViTConvLayer(
            config, out_channels=in_channels, kernel_size=config.conv_kernel_size, name="fusion"
        )