src/transformers/models/mobilevit/modeling_tf_mobilevit.py/TFMobileViTMainLayer/__init__
class TFMobileViTMainLayer:    def __init__(self, config: MobileViTConfig, expand_output: bool = True, **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.expand_output = expand_output

        self.conv_stem = TFMobileViTConvLayer(
            config,
            out_channels=config.neck_hidden_sizes[0],
            kernel_size=3,
            stride=2,
            name="conv_stem",
        )

        self.encoder = TFMobileViTEncoder(config, name="encoder")

        if self.expand_output:
            self.conv_1x1_exp = TFMobileViTConvLayer(
                config, out_channels=config.neck_hidden_sizes[6], kernel_size=1, name="conv_1x1_exp"
            )

        self.pooler = tf.keras.layers.GlobalAveragePooling2D(data_format="channels_first", name="pooler")