codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the tensorflow segformer model from future import annotations import inspect import unittest from typing import list tuple from transformers import segformerconfig from transformers fileutils import istfavailable isvisionavailable from transformers testingutils import requiretf slow from testconfigurationcommon import configtester from testmodelingtfcommon import tfmodeltestermixin floatstensor idstensor from testpipelinemixin import pipelinetestermixin if istfavailable import numpy as np import tensorflow as tf from transformers import tfsegformerforimageclassification tfsegformerforsemanticsegmentation tfsegformermodel from transformers models segformer modelingtfsegformer import tfsegformerpretrainedmodelarchivelist if isvisionavailable from pil import image from transformers import segformerimageprocessor class tfsegformerconfigtesterconfigtester def createandtestconfigcommonpropertiesself config self configclassself inputsdict self parent asserttruehasattrconfig hiddensizes self parent asserttruehasattrconfig numattentionheads self parent asserttruehasattrconfig numencoderblocks class tfsegformermodeltester def init self parent batchsize13 imagesize64 numchannels3 numencoderblocks4 depths1 1 1 1 srratios8 4 2 1 hiddensizes8 8 16 16 downsamplingrates1 4 8 16 numattentionheads1 1 2 2 istrainingtrue uselabelstrue hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 initializerrange0 02 numlabels3 scopenone self parent parent self batchsize batchsize self imagesize imagesize self numchannels numchannels self numencoderblocks numencoderblocks self srratios srratios self depths depths self hiddensizes hiddensizes self downsamplingrates downsamplingrates self numattentionheads numattentionheads self istraining istraining self uselabels uselabels self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self initializerrange initializerrange self numlabels numlabels self scope scope def prepareconfigandinputsself pixelvalues floatstensorself batchsize self numchannels self imagesize self imagesize labels none if self uselabels labels idstensorself batchsize self imagesize self imagesize self numlabels config self getconfig return config pixelvalues labels def getconfigself return segformerconfig imagesizeself imagesize numchannelsself numchannels numencoderblocksself numencoderblocks depthsself depths hiddensizesself hiddensizes numattentionheadsself numattentionheads hiddenactself hiddenact hiddendropoutprobself hiddendropoutprob attentionprobsdropoutprobself attentionprobsdropoutprob initializerrangeself initializerrange numlabelsself numlabels def createandcheckmodelself config pixelvalues labels model tfsegformermodelconfigconfig result modelpixelvalues trainingfalse expectedheight expectedwidth self imagesize self downsamplingrates1 2 self parent assertequal result lasthiddenstate shape self batchsize self hiddensizes1 expectedheight expectedwidth def createandcheckforimagesegmentationself config pixelvalues labels config numlabels self numlabels model tfsegformerforsemanticsegmentationconfig result modelpixelvalues trainingfalse self parent assertequal result logits shape self batchsize self numlabels self imagesize 4 self imagesize 4 result modelpixelvalues labelslabels trainingfalse self parent assertequal result logits shape self batchsize self numlabels self imagesize 4 self imagesize 4 def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config pixelvalues labels configandinputs inputsdict pixelvalues pixelvalues return config inputsdict def prepareconfigandinputsforkerasfitself forsegmentation bool false configandinputs self prepareconfigandinputs config pixelvalues seglabels configandinputs if forsegmentation inputsdict pixelvalues pixelvalues labels seglabels else inputsdict pixelvalues pixelvalues labels tf zerosself batchsize return config inputsdict requiretf class tfsegformermodeltesttfmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses tfsegformermodel tfsegformerforimageclassification tfsegformerforsemanticsegmentation if istfavailable else pipelinemodelmapping featureextraction tfsegformermodel imageclassification tfsegformerforimageclassification if istfavailable else testheadmasking false testonnx false testpruning false testresizeembeddings false def setupself self modeltester tfsegformermodeltesterself self configtester tfsegformerconfigtesterself configclasssegformerconfig hastextmodalityfalse def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs unittest skipsegformer does not use inputsembeds def testinputsembedsself pass unittest skipsegformer does not have getinputembeddings method and getoutputembeddings methods def testmodelcommonattributesself pass def testforwardsignatureself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig signature inspect signaturemodel call signature parameters is an ordereddict so argnames order is deterministic argnames signature parameters keys expectedargnames pixelvalues self assertlistequalargnames 1 expectedargnames def testattentionoutputsself config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true for modelclass in self allmodelclasses inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions expectednumattentions sumself modeltester depths self assertequallenattentions expectednumattentions check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions expectednumattentions verify the first attentions first block first layer expectedseqlen self modeltester imagesize 4 2 expectedreducedseqlen self modeltester imagesize 4 self modeltester srratios0 2 self assertlistequal listattentions0 shape3 self modeltester numattentionheads0 expectedseqlen expectedreducedseqlen verify the last attentions last block last layer expectedseqlen self modeltester imagesize 32 2 expectedreducedseqlen self modeltester imagesize 32 self modeltester srratios1 2 self assertlistequal listattentions1 shape3 self modeltester numattentionheads1 expectedseqlen expectedreducedseqlen outlen lenoutputs check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig outputs modelself prepareforclassinputsdict modelclass self assertequaloutlen 1 lenoutputs selfattentions outputs attentions self assertequallenselfattentions expectednumattentions verify the first attentions first block first layer expectedseqlen self modeltester imagesize 4 2 expectedreducedseqlen self modeltester imagesize 4 self modeltester srratios0 2 self assertlistequal listselfattentions0 shape3 self modeltester numattentionheads0 expectedseqlen expectedreducedseqlen def testhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig outputs modelself prepareforclassinputsdict modelclass hiddenstates outputs hiddenstates expectednumlayers self modeltester numencoderblocks self assertequallenhiddenstates expectednumlayers verify the first hidden states first block self assertlistequal listhiddenstates0 shape3 self modeltester hiddensizes0 self modeltester imagesize 4 self modeltester imagesize 4 config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass def testmodeloutputsequivalenceself config inputsdict self modeltester prepareconfigandinputsforcommon def checkequivalencemodel tupleinputs dictinputs additionalkwargs tupleoutput modeltupleinputs returndictfalse additionalkwargs dictoutput modeldictinputs returndicttrue additionalkwargs totuple def recursivechecktupleobject dictobject if isinstancetupleobject list tuple for tupleiterablevalue dictiterablevalue in ziptupleobject dictobject recursivechecktupleiterablevalue dictiterablevalue elif tupleobject is none return else self asserttrue alltf equaltupleobject dictobject msg tuple and dict output are not equal difference f tf math reducemaxtf abstupleobject dictobject recursivechecktupleoutput dictoutput for modelclass in self allmodelclasses model modelclassconfig tupleinputs self prepareforclassinputsdict modelclass dictinputs self prepareforclassinputsdict modelclass checkequivalencemodel tupleinputs dictinputs tupleinputs self prepareforclassinputsdict modelclass dictinputs self prepareforclassinputsdict modelclass checkequivalencemodel tupleinputs dictinputs outputhiddenstates true if self hasattentions tupleinputs self prepareforclassinputsdict modelclass dictinputs self prepareforclassinputsdict modelclass checkequivalencemodel tupleinputs dictinputs outputattentions true todo incorporate label support for semantic segmentation in testmodelingtfcommon py unittest skipif not istfavailable or lentf config listphysicaldevicesgpu 0 reasontf does not support backprop for grouped convolutions on cpu def testdatasetconversionself super testdatasetconversion def checkkerasfitresultsself valloss1 valloss2 atol2e1 rtol2e1 self asserttruenp allclosevalloss1 valloss2 atolatol rtolrtol unittest skipif not istfavailable or lentf config listphysicaldevicesgpu 0 reasontf does not support backprop for grouped convolutions on cpu slow def testkerasfitself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses since tfsegformermodel cannot operate with the default fit method if modelclass name tfsegformermodel model modelclassconfig if getattrmodel hfcomputeloss none super testkerasfit def testlosscomputationself config inputsdict self modeltester prepareconfigandinputsforcommon def applymodel forsegmentation true if modelclass name tfsegformerforsemanticsegmentation else false the number of elements in the loss should be the same as the number of elements in the label preparedforclass self modeltester prepareconfigandinputsforkerasfit forsegmentationforsegmentation addedlabel preparedforclasssortedpreparedforclass keys inputsdict keys reversetrue0 losssize tf sizeaddedlabel test that model correctly compute the loss with kwargs possibleinputnames inputids pixelvalues inputfeatures inputname possibleinputnames intersectionsetpreparedforclass pop modelinput preparedforclass popinputname loss modelmodelinput preparedforclass0 if modelclass name tfsegformerforsemanticsegmentation semantic segmentation loss is computed similarly as https github comhuggingfacetransformersblobmainsrctransformersmodelingtfutils pyl210 self assertequalloss shape 1 else self assertequalloss shape losssize test that model correctly compute the loss with a dict preparedforclass self modeltester prepareconfigandinputsforkerasfit forsegmentationforsegmentation loss modelpreparedforclass0 if modelclass name tfsegformerforsemanticsegmentation self assertequalloss shape 1 else self assertequalloss shape losssize test that model correctly compute the loss with a tuple labelkeys preparedforclass keys inputsdict keys signature inspect signaturemodel call parameters signaturenames listsignature keys create a dictionary holding the location of the tensors in the tuple tupleindexmapping 0 inputname for labelkey in labelkeys labelkeyindex signaturenames indexlabelkey tupleindexmappinglabelkeyindex labelkey sortedtupleindexmapping sortedtupleindexmapping items initialize a list with their default values update the values and convert to a tuple listinput for name in signaturenames if name kwargs listinput appendsignaturename default for index value in sortedtupleindexmapping listinputindex preparedforclassvalue tupleinput tuplelistinput send to model loss modeltupleinput 10 if modelclass name tfsegformerforsemanticsegmentation self assertequalloss shape 1 else self assertequalloss shape losssize for modelclass in self allmodelclasses since tfsegformermodel won t have labels against which we could compute loss if modelclass name tfsegformermodel model modelclassconfig applymodel def checkpttfoutputsself tfoutputs ptoutputs modelclass tol2e4 nameoutputs attributesnone we override with a slightly higher tol value as semseg models tend to diverge a bit more super checkpttfoutputstfoutputs ptoutputs modelclass tol name attributes slow def testmodelfrompretrainedself for modelname in tfsegformerpretrainedmodelarchivelist 1 model tfsegformermodel frompretrainedmodelname self assertisnotnonemodel we will verify our results on an image of cute cats def prepareimg image image open testsfixturestestssamplescoco000000039769 png return image requiretf class tfsegformermodelintegrationtestunittest testcase slow def testinferenceimagesegmentationadeself only resize normalize imageprocessor segformerimageprocessor imagescale512 512 keepratiofalse alignfalse dorandomcropfalse model tfsegformerforsemanticsegmentation frompretrainednvidiasegformerb0finetunedade512512 image prepareimg encodedinputs imageprocessorimagesimage returntensorstf pixelvalues encodedinputs pixelvalues outputs modelpixelvalues trainingfalse expectedshape tf tensorshape1 model config numlabels 128 128 self assertequaloutputs logits shape expectedshape expectedslice tf constant 4 6310 5 5232 6 2356 5 1921 6 1444 6 5996 5 4424 6 2790 6 7574 12 1391 13 3122 13 9554 12 8732 13 9352 14 3563 12 9438 13 8226 14 2513 12 5134 13 4686 14 4915 12 8669 14 4343 14 7758 13 2523 14 5819 15 0694 tf debugging assertnearoutputs logits0 3 3 3 expectedslice atol1e4 slow def testinferenceimagesegmentationcityself only resize normalize imageprocessor segformerimageprocessor imagescale512 512 keepratiofalse alignfalse dorandomcropfalse model tfsegformerforsemanticsegmentation frompretrained nvidiasegformerb1finetunedcityscapes10241024 image prepareimg encodedinputs imageprocessorimagesimage returntensorstf pixelvalues encodedinputs pixelvalues outputs modelpixelvalues trainingfalse expectedshape tf tensorshape1 model config numlabels 128 128 self assertequaloutputs logits shape expectedshape expectedslice tf constant 13 5748 13 9111 12 6500 14 3500 15 3683 14 2328 14 7532 16 0424 15 6087 17 1651 15 8725 12 9653 17 2580 17 3718 14 8223 16 6058 16 8783 16 7452 3 6456 3 0209 1 4203 3 0797 3 1959 2 0000 1 8757 1 9217 1 6997 tf debugging assertnearoutputs logits0 3 3 3 expectedslice atol1e1 coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the tensorflow segformer model signature parameters is an ordereddict so arg_names order is deterministic check that output_attentions also work using config verify the first attentions first block first layer verify the last attentions last block last layer check attention is always last and order is fine verify the first attentions first block first layer verify the first hidden states first block check that output_hidden_states also work using config todo incorporate label support for semantic segmentation in test_modeling_tf_common py since tfsegformermodel cannot operate with the default fit method the number of elements in the loss should be the same as the number of elements in the label test that model correctly compute the loss with kwargs semantic segmentation loss is computed similarly as https github com huggingface transformers blob main src transformers modeling_tf_utils py l210 test that model correctly compute the loss with a dict test that model correctly compute the loss with a tuple create a dictionary holding the location of the tensors in the tuple initialize a list with their default values update the values and convert to a tuple send to model since tfsegformermodel won t have labels against which we could compute loss we override with a slightly higher tol value as semseg models tend to diverge a bit more we will verify our results on an image of cute cats only resize normalize only resize normalize
from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import SegformerConfig
from transformers.file_utils import is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import numpy as np
    import tensorflow as tf

    from transformers import TFSegformerForImageClassification, TFSegformerForSemanticSegmentation, TFSegformerModel
    from transformers.models.segformer.modeling_tf_segformer import TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
    from PIL import Image

    from transformers import SegformerImageProcessor


class TFSegformerConfigTester(ConfigTester):
    def create_and_test_config_common_properties(self):
        config = self.config_class(**self.inputs_dict)
        self.parent.assertTrue(hasattr(config, "hidden_sizes"))
        self.parent.assertTrue(hasattr(config, "num_attention_heads"))
        self.parent.assertTrue(hasattr(config, "num_encoder_blocks"))


class TFSegformerModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        image_size=64,
        num_channels=3,
        num_encoder_blocks=4,
        depths=[1, 1, 1, 1],
        sr_ratios=[8, 4, 2, 1],
        hidden_sizes=[8, 8, 16, 16],
        downsampling_rates=[1, 4, 8, 16],
        num_attention_heads=[1, 1, 2, 2],
        is_training=True,
        use_labels=True,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        initializer_range=0.02,
        num_labels=3,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.image_size = image_size
        self.num_channels = num_channels
        self.num_encoder_blocks = num_encoder_blocks
        self.sr_ratios = sr_ratios
        self.depths = depths
        self.hidden_sizes = hidden_sizes
        self.downsampling_rates = downsampling_rates
        self.num_attention_heads = num_attention_heads
        self.is_training = is_training
        self.use_labels = use_labels
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.initializer_range = initializer_range
        self.num_labels = num_labels
        self.scope = scope

    def prepare_config_and_inputs(self):
        pixel_values = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

        labels = None
        if self.use_labels:
            labels = ids_tensor([self.batch_size, self.image_size, self.image_size], self.num_labels)

        config = self.get_config()
        return config, pixel_values, labels

    def get_config(self):
        return SegformerConfig(
            image_size=self.image_size,
            num_channels=self.num_channels,
            num_encoder_blocks=self.num_encoder_blocks,
            depths=self.depths,
            hidden_sizes=self.hidden_sizes,
            num_attention_heads=self.num_attention_heads,
            hidden_act=self.hidden_act,
            hidden_dropout_prob=self.hidden_dropout_prob,
            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
            initializer_range=self.initializer_range,
            num_labels=self.num_labels,
        )

    def create_and_check_model(self, config, pixel_values, labels):
        model = TFSegformerModel(config=config)
        result = model(pixel_values, training=False)
        expected_height = expected_width = self.image_size // (self.downsampling_rates[-1] * 2)
        self.parent.assertEqual(
            result.last_hidden_state.shape, (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)
        )

    def create_and_check_for_image_segmentation(self, config, pixel_values, labels):
        config.num_labels = self.num_labels
        model = TFSegformerForSemanticSegmentation(config)
        result = model(pixel_values, training=False)
        self.parent.assertEqual(
            result.logits.shape, (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)
        )
        result = model(pixel_values, labels=labels, training=False)
        self.parent.assertEqual(
            result.logits.shape, (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)
        )

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, pixel_values, labels = config_and_inputs
        inputs_dict = {"pixel_values": pixel_values}
        return config, inputs_dict

    def prepare_config_and_inputs_for_keras_fit(self, for_segmentation: bool = False):
        config_and_inputs = self.prepare_config_and_inputs()
        config, pixel_values, seg_labels = config_and_inputs
        if for_segmentation:
            inputs_dict = {"pixel_values": pixel_values, "labels": seg_labels}
        else:
            inputs_dict = {"pixel_values": pixel_values, "labels": tf.zeros((self.batch_size))}
        return config, inputs_dict


@require_tf
class TFSegformerModelTest(TFModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    all_model_classes = (
        (TFSegformerModel, TFSegformerForImageClassification, TFSegformerForSemanticSegmentation)
        if is_tf_available()
        else ()
    )
    pipeline_model_mapping = (
        {"feature-extraction": TFSegformerModel, "image-classification": TFSegformerForImageClassification}
        if is_tf_available()
        else {}
    )

    test_head_masking = False
    test_onnx = False
    test_pruning = False
    test_resize_embeddings = False

    def setUp(self):
        self.model_tester = TFSegformerModelTester(self)
        self.config_tester = TFSegformerConfigTester(self, config_class=SegformerConfig, has_text_modality=False)

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    @unittest.skip("SegFormer does not use inputs_embeds")
    def test_inputs_embeds(self):
        pass

    @unittest.skip("SegFormer does not have get_input_embeddings method and get_output_embeddings methods")
    def test_model_common_attributes(self):
        pass

    def test_forward_signature(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            signature = inspect.signature(model.call)
            
            arg_names = [*signature.parameters.keys()]

            expected_arg_names = ["pixel_values"]
            self.assertListEqual(arg_names[:1], expected_arg_names)

    def test_attention_outputs(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions

            expected_num_attentions = sum(self.model_tester.depths)
            self.assertEqual(len(attentions), expected_num_attentions)

            
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class))
            attentions = outputs.attentions

            self.assertEqual(len(attentions), expected_num_attentions)

            
            expected_seq_len = (self.model_tester.image_size // 4) ** 2
            expected_reduced_seq_len = (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
            self.assertListEqual(
                list(attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len],
            )

            
            expected_seq_len = (self.model_tester.image_size // 32) ** 2
            expected_reduced_seq_len = (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
            self.assertListEqual(
                list(attentions[-1].shape[-3:]),
                [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len],
            )
            out_len = len(outputs)

            
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            self.assertEqual(out_len + 1, len(outputs))

            self_attentions = outputs.attentions

            self.assertEqual(len(self_attentions), expected_num_attentions)
            
            expected_seq_len = (self.model_tester.image_size // 4) ** 2
            expected_reduced_seq_len = (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]),
                [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len],
            )

    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)

            outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            hidden_states = outputs.hidden_states

            expected_num_layers = self.model_tester.num_encoder_blocks
            self.assertEqual(len(hidden_states), expected_num_layers)

            
            self.assertListEqual(
                list(hidden_states[0].shape[-3:]),
                [
                    self.model_tester.hidden_sizes[0],
                    self.model_tester.image_size // 4,
                    self.model_tester.image_size // 4,
                ],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)

    def test_model_outputs_equivalence(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}):
            tuple_output = model(tuple_inputs, return_dict=False, **additional_kwargs)
            dict_output = model(dict_inputs, return_dict=True, **additional_kwargs).to_tuple()

            def recursive_check(tuple_object, dict_object):
                if isinstance(tuple_object, (List, Tuple)):
                    for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object):
                        recursive_check(tuple_iterable_value, dict_iterable_value)
                elif tuple_object is None:
                    return
                else:
                    self.assertTrue(
                        all(tf.equal(tuple_object, dict_object)),
                        msg=(
                            "Tuple and dict output are not equal. Difference:"
                            f" {tf.math.reduce_max(tf.abs(tuple_object - dict_object))}"
                        ),
                    )

                recursive_check(tuple_output, dict_output)

        for model_class in self.all_model_classes:
            model = model_class(config)

            tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
            dict_inputs = self._prepare_for_class(inputs_dict, model_class)
            check_equivalence(model, tuple_inputs, dict_inputs)

            tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
            dict_inputs = self._prepare_for_class(inputs_dict, model_class)
            check_equivalence(model, tuple_inputs, dict_inputs, {"output_hidden_states": True})

            if self.has_attentions:
                tuple_inputs = self._prepare_for_class(inputs_dict, model_class)
                dict_inputs = self._prepare_for_class(inputs_dict, model_class)
                check_equivalence(model, tuple_inputs, dict_inputs, {"output_attentions": True})

            

    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices("GPU")) == 0,
        reason="TF does not support backprop for grouped convolutions on CPU.",
    )
    def test_dataset_conversion(self):
        super().test_dataset_conversion()

    def check_keras_fit_results(self, val_loss1, val_loss2, atol=2e-1, rtol=2e-1):
        self.assertTrue(np.allclose(val_loss1, val_loss2, atol=atol, rtol=rtol))

    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices("GPU")) == 0,
        reason="TF does not support backprop for grouped convolutions on CPU.",
    )
    @slow
    def test_keras_fit(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            
            if model_class.__name__ != "TFSegformerModel":
                model = model_class(config)
                if getattr(model, "hf_compute_loss", None):
                    super().test_keras_fit()

    def test_loss_computation(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        def apply(model):
            for_segmentation = True if model_class.__name__ == "TFSegformerForSemanticSegmentation" else False
            
            _, prepared_for_class = self.model_tester.prepare_config_and_inputs_for_keras_fit(
                for_segmentation=for_segmentation
            )
            added_label = prepared_for_class[sorted(prepared_for_class.keys() - inputs_dict.keys(), reverse=True)[0]]
            loss_size = tf.size(added_label)

            
            possible_input_names = {"input_ids", "pixel_values", "input_features"}
            input_name = possible_input_names.intersection(set(prepared_for_class)).pop()
            model_input = prepared_for_class.pop(input_name)

            loss = model(model_input, **prepared_for_class)[0]

            if model_class.__name__ == "TFSegformerForSemanticSegmentation":
                
                
                self.assertEqual(loss.shape, (1,))
            else:
                self.assertEqual(loss.shape, [loss_size])

            
            _, prepared_for_class = self.model_tester.prepare_config_and_inputs_for_keras_fit(
                for_segmentation=for_segmentation
            )
            loss = model(**prepared_for_class)[0]

            if model_class.__name__ == "TFSegformerForSemanticSegmentation":
                self.assertEqual(loss.shape, (1,))
            else:
                self.assertEqual(loss.shape, [loss_size])

            
            label_keys = prepared_for_class.keys() - inputs_dict.keys()
            signature = inspect.signature(model.call).parameters
            signature_names = list(signature.keys())

            
            tuple_index_mapping = {0: input_name}
            for label_key in label_keys:
                label_key_index = signature_names.index(label_key)
                tuple_index_mapping[label_key_index] = label_key
            sorted_tuple_index_mapping = sorted(tuple_index_mapping.items())
            
            list_input = []

            for name in signature_names:
                if name != "kwargs":
                    list_input.append(signature[name].default)

            for index, value in sorted_tuple_index_mapping:
                list_input[index] = prepared_for_class[value]

            tuple_input = tuple(list_input)

            
            loss = model(tuple_input[:-1])[0]
            if model_class.__name__ == "TFSegformerForSemanticSegmentation":
                self.assertEqual(loss.shape, (1,))
            else:
                self.assertEqual(loss.shape, [loss_size])

        for model_class in self.all_model_classes:
            
            
            if model_class.__name__ != "TFSegformerModel":
                model = model_class(config)
                apply(model)

    def check_pt_tf_outputs(self, tf_outputs, pt_outputs, model_class, tol=2e-4, name="outputs", attributes=None):
        
        super().check_pt_tf_outputs(tf_outputs, pt_outputs, model_class, tol, name, attributes)

    @slow
    def test_model_from_pretrained(self):
        for model_name in TF_SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = TFSegformerModel.from_pretrained(model_name)
            self.assertIsNotNone(model)



def prepare_img():
    image = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
    return image


@require_tf
class TFSegformerModelIntegrationTest(unittest.TestCase):
    @slow
    def test_inference_image_segmentation_ade(self):
        
        image_processor = SegformerImageProcessor(
            image_scale=(512, 512), keep_ratio=False, align=False, do_random_crop=False
        )
        model = TFSegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")

        image = prepare_img()
        encoded_inputs = image_processor(images=image, return_tensors="tf")
        pixel_values = encoded_inputs.pixel_values

        outputs = model(pixel_values, training=False)

        expected_shape = tf.TensorShape((1, model.config.num_labels, 128, 128))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = tf.constant(
            [
                [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
                [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
                [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
            ]
        )
        tf.debugging.assert_near(outputs.logits[0, :3, :3, :3], expected_slice, atol=1e-4)

    @slow
    def test_inference_image_segmentation_city(self):
        
        image_processor = SegformerImageProcessor(
            image_scale=(512, 512), keep_ratio=False, align=False, do_random_crop=False
        )
        model = TFSegformerForSemanticSegmentation.from_pretrained(
            "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
        )

        image = prepare_img()
        encoded_inputs = image_processor(images=image, return_tensors="tf")
        pixel_values = encoded_inputs.pixel_values

        outputs = model(pixel_values, training=False)

        expected_shape = tf.TensorShape((1, model.config.num_labels, 128, 128))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = tf.constant(
            [
                [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
                [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
                [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
            ]
        )
        tf.debugging.assert_near(outputs.logits[0, :3, :3, :3], expected_slice, atol=1e-1)