src/transformers/models/data2vec/modeling_tf_data2vec_vision.py/TFAdaptiveAvgPool2D/call
class TFAdaptiveAvgPool2D:    def call(self, inputs):
        # Rearrange from NHWC -> NCHW
        inputs = tf.transpose(inputs, perm=[0, 3, 1, 2])
        # Perform W-pooling
        inputs = self.w_pool(inputs)
        # Rearrange NCHW -> NCWH
        inputs = tf.transpose(inputs, perm=[0, 1, 3, 2])
        # Perform H-pooling
        inputs = self.h_pool(inputs)
        # Rearrange from NCWH -> NHWC
        inputs = tf.transpose(inputs, perm=[0, 3, 2, 1])
        return inputs