



import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
    import torch

    from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin

class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				Optional[Any]    ,      a__	:				List[str]    ,      a__	:				Optional[int]=None    ,      a__	:				Any=None    ,      a__	:				str=None    ,      a__	:				str="resnet50"    ,      a__	:				Any=3    ,      a__	:				Union[str, Any]=32    ,      a__	:				Optional[Any]=3    ,      a__	:				Optional[int]=True    ,      a__	:				List[Any]=True    ,      ):





         """simple docstring"""



         __snake_case								=       parent
         __snake_case								=       out_indices if out_indices is not None else [4]
         __snake_case								=       stage_names
         __snake_case								=       out_features
         __snake_case								=       backbone
         __snake_case								=       batch_size
         __snake_case								=       image_size
         __snake_case								=       num_channels
         __snake_case								=       use_pretrained_backbone
         __snake_case								=       is_training


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)
         __snake_case								=       self.get_config()

         return config, pixel_values


      def        a			(self	:				Any				):





         """simple docstring"""



         return TimmBackboneConfig(
             image_size=self.image_size    ,      num_channels=self.num_channels    ,      out_features=self.out_features    ,      out_indices=self.out_indices    ,      stage_names=self.stage_names    ,      use_pretrained_backbone=self.use_pretrained_backbone    ,      backbone=self.backbone    ,      )


      def        a			(self	:				str    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[int]				):





         """simple docstring"""



         __snake_case								=       TimmBackbone(config=a__				)
         model.to(a__				)
         model.eval()
         with torch.no_grad():
            __snake_case								=       model(a__				)
         self.parent.assertEqual(
             result.feature_map[-1].shape    ,      (self.batch_size, model.channels[-1], 14, 14)    ,      )


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       self.prepare_config_and_inputs()
         __snake_case	,     __snake_case								=       config_and_inputs
         __snake_case								=       {'''pixel_values''': pixel_values}
         return config, inputs_dict






@require_torch
@require_timm
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Optional[Any]							=   (TimmBackbone,) if is_torch_available() else ()
      A_						:    List[Any]							=   {'feature-extraction': TimmBackbone} if is_torch_available() else {}
      A_						:    Dict							=   False
      A_						:    List[str]							=   False
      A_						:    Optional[int]							=   False
      A_						:    Tuple							=   False


      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         __snake_case								=       TimmBackboneModelTester(self				)
         __snake_case								=       ConfigTester(self    ,      config_class=a__    ,      has_text_modality=a__				)


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         self.config_tester.create_and_test_config_to_json_string()
         self.config_tester.create_and_test_config_to_json_file()
         self.config_tester.create_and_test_config_from_and_save_pretrained()
         self.config_tester.create_and_test_config_with_num_labels()
         self.config_tester.check_config_can_be_init_without_params()
         self.config_tester.check_config_arguments_init()


      def        a			(self	:				str				):





         """simple docstring"""



         __snake_case								=       '''resnet18'''
         __snake_case								=       '''microsoft/resnet-18'''

         __snake_case								=       AutoBackbone.from_pretrained(a__    ,      use_timm_backbone=a__				)
         __snake_case								=       AutoBackbone.from_pretrained(a__				)

         self.assertEqual(len(timm_model.out_features				)    ,      len(transformers_model.out_features				)				)
         self.assertEqual(len(timm_model.stage_names				)    ,      len(transformers_model.stage_names				)				)
         self.assertEqual(timm_model.channels    ,      transformers_model.channels				)
         # Out indices are set to the last layer by default. For timm models, we don't know
         # the number of layers in advance, so we set it to (-1,), whereas for transformers
         # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
         self.assertEqual(timm_model.out_indices    ,      (-1,)				)
         self.assertEqual(transformers_model.out_indices    ,      [len(timm_model.stage_names				) - 1]				)

         __snake_case								=       AutoBackbone.from_pretrained(a__    ,      use_timm_backbone=a__    ,      out_indices=[1, 2, 3]				)
         __snake_case								=       AutoBackbone.from_pretrained(a__    ,      out_indices=[1, 2, 3]				)

         self.assertEqual(timm_model.out_indices    ,      transformers_model.out_indices				)
         self.assertEqual(len(timm_model.out_features				)    ,      len(transformers_model.out_features				)				)
         self.assertEqual(timm_model.channels    ,      transformers_model.channels				)


      @unittest.skip('''TimmBackbone doesn\'t support feed forward chunking'''				)
      def        a			(self	:				List[str]				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone doesn\'t have num_hidden_layers attribute'''				)
      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone initialization is managed on the timm side'''				)
      def        a			(self	:				str				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
      def        a			(self	:				List[str]				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
      def        a			(self	:				Dict				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone model cannot be created without specifying a backbone checkpoint'''				)
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         pass


      @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         pass


      @unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         pass


      @unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         pass


      @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         pass


      @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
      def        a			(self	:				Any				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone doesn\'t have hidden size info in its configuration.'''				)
      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         pass


      @unittest.skip('''TimmBackbone doesn\'t support output_attentions.'''				)
      def        a			(self	:				Any				):





         """simple docstring"""



         pass


      @unittest.skip('''Safetensors is not supported by timm.'''				)
      def        a			(self	:				Any				):





         """simple docstring"""



         pass


      @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''				)
      def        a			(self	:				int				):





         """simple docstring"""



         pass


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            __snake_case								=       model_class(a__				)
            __snake_case								=       inspect.signature(model.forward				)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            __snake_case								=       [*signature.parameters.keys()]

            __snake_case								=       ['''pixel_values''']
            self.assertListEqual(arg_names[:1]    ,      a__				)


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case								=       True
         __snake_case								=       self.has_attentions

         # no need to test all models as different heads yield the same functionality
         __snake_case								=       self.all_model_classes[0]
         __snake_case								=       model_class(a__				)
         model.to(a__				)

         __snake_case								=       self._prepare_for_class(a__    ,      a__				)
         __snake_case								=       model(**a__				)
         __snake_case								=       outputs[0][-1]

         # Encoder-/Decoder-only models
         __snake_case								=       outputs.hidden_states[0]
         hidden_states.retain_grad()

         if self.has_attentions:
            __snake_case								=       outputs.attentions[0]
            attentions.retain_grad()

         output.flatten()[0].backward(retain_graph=a__				)

         self.assertIsNotNone(hidden_states.grad				)

         if self.has_attentions:
            self.assertIsNotNone(attentions.grad				)


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            __snake_case								=       model_class(a__				)
            model.to(a__				)
            model.eval()
            __snake_case								=       model(**a__				)

            self.assertEqual(len(result.feature_maps				)    ,      len(config.out_indices				)				)
            self.assertEqual(len(model.channels				)    ,      len(config.out_indices				)				)

            # Check output of last stage is taken if out_features=None, out_indices=None
            __snake_case								=       copy.deepcopy(a__				)
            __snake_case								=       None
            __snake_case								=       model_class(a__				)
            model.to(a__				)
            model.eval()
            __snake_case								=       model(**a__				)

            self.assertEqual(len(result.feature_maps				)    ,      1				)
            self.assertEqual(len(model.channels				)    ,      1				)

            # Check backbone can be initialized with fresh weights
            __snake_case								=       copy.deepcopy(a__				)
            __snake_case								=       False
            __snake_case								=       model_class(a__				)
            model.to(a__				)
            model.eval()
            __snake_case								=       model(**a__				)








class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				str    ,      a__	:				list				):





         """simple docstring"""



         __snake_case								=       set_counts
         __snake_case								=       max(a__				)
         __snake_case								=       len(a__				)
         __snake_case								=       [1] * num_sets
         __snake_case								=       list(range(a__				)				)


      def        a			(self	:				str    ,      a__	:				int    ,      a__	:				int				):





         """simple docstring"""



         __snake_case								=       self.get_parent(a__				)
         __snake_case								=       self.get_parent(a__				)

         if src_parent == dst_parent:
            return False

         if self.ranks[dst_parent] >= self.ranks[src_parent]:
            self.set_counts[dst_parent] += self.set_counts[src_parent]
            __snake_case								=       0
            __snake_case								=       dst_parent
            if self.ranks[dst_parent] == self.ranks[src_parent]:
               self.ranks[dst_parent] += 1
            __snake_case								=       self.set_counts[dst_parent]
         else:
            self.set_counts[src_parent] += self.set_counts[dst_parent]
            __snake_case								=       0
            __snake_case								=       src_parent
            __snake_case								=       self.set_counts[src_parent]

         __snake_case								=       max(self.max_set    ,      a__				)
         return True


      def        a			(self	:				Union[str, Any]    ,      a__	:				int				):





         """simple docstring"""



         if self.parents[disj_set] == disj_set:
            return disj_set
         __snake_case								=       self.get_parent(self.parents[disj_set]				)
         return self.parents[disj_set]

