src/transformers/models/convnext/modeling_tf_convnext.py/TFConvNextStage/__init__
class TFConvNextStage:    def __init__(
        self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None, **kwargs
    ):
        super().__init__(**kwargs)
        if in_channels != out_channels or stride > 1:
            self.downsampling_layer = [
                tf.keras.layers.LayerNormalization(
                    epsilon=1e-6,
                    name="downsampling_layer.0",
                ),
                # Inputs to this layer will follow NHWC format since we
                # transposed the inputs from NCHW to NHWC in the `TFConvNextEmbeddings`
                # layer. All the outputs throughout the model will be in NHWC
                # from this point on until the output where we again change to
                # NCHW.
                tf.keras.layers.Conv2D(
                    filters=out_channels,
                    kernel_size=kernel_size,
                    strides=stride,
                    kernel_initializer=get_initializer(config.initializer_range),
                    bias_initializer="zeros",
                    name="downsampling_layer.1",
                ),
            ]
        else:
            self.downsampling_layer = [tf.identity]

        drop_path_rates = drop_path_rates or [0.0] * depth
        self.layers = [
            TFConvNextLayer(
                config,
                dim=out_channels,
                drop_path=drop_path_rates[j],
                name=f"layers.{j}",
            )
            for j in range(depth)
        ]