"""simple docstring"""


import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
        MODEL_MAPPING,
        LevitForImageClassification,
        LevitForImageClassificationWithTeacher,
        LevitModel,
    )
    from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import LevitImageProcessor


class 				a (						lowerCAmelCase_      ):


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      _UpperCAmelCase         =					self.config_class(**self.inputs_dict    )
      self.parent.assertTrue(hasattr(__lowerCAmelCase      , """hidden_sizes"""    )    )
      self.parent.assertTrue(hasattr(__lowerCAmelCase      , """num_attention_heads"""    )    )



class 				a :


    def __init__(    self		:     Dict      , __lowerCAmelCase		:     Dict      , __lowerCAmelCase		:     List[str]=13      , __lowerCAmelCase		:     Tuple=64      , __lowerCAmelCase		:     List[str]=3      , __lowerCAmelCase		:     Tuple=3      , __lowerCAmelCase		:     Dict=2      , __lowerCAmelCase		:     List[Any]=1      , __lowerCAmelCase		:     Any=16      , __lowerCAmelCase		:     Any=[128, 256, 384]      , __lowerCAmelCase		:     Optional[Any]=[4, 6, 8]      , __lowerCAmelCase		:     Optional[Any]=[2, 3, 4]      , __lowerCAmelCase		:     Tuple=[16, 16, 16]      , __lowerCAmelCase		:     Any=0      , __lowerCAmelCase		:     Optional[Any]=[2, 2, 2]      , __lowerCAmelCase		:     Optional[Any]=[2, 2, 2]      , __lowerCAmelCase		:     Any=0.02      , __lowerCAmelCase		:     int=True      , __lowerCAmelCase		:     Union[str, Any]=True      , __lowerCAmelCase		:     List[Any]=2      , ):
      _UpperCAmelCase         =					parent
      _UpperCAmelCase         =					batch_size
      _UpperCAmelCase         =					image_size
      _UpperCAmelCase         =					num_channels
      _UpperCAmelCase         =					kernel_size
      _UpperCAmelCase         =					stride
      _UpperCAmelCase         =					padding
      _UpperCAmelCase         =					hidden_sizes
      _UpperCAmelCase         =					num_attention_heads
      _UpperCAmelCase         =					depths
      _UpperCAmelCase         =					key_dim
      _UpperCAmelCase         =					drop_path_rate
      _UpperCAmelCase         =					patch_size
      _UpperCAmelCase         =					attention_ratio
      _UpperCAmelCase         =					mlp_ratio
      _UpperCAmelCase         =					initializer_range
      _UpperCAmelCase         =					[
          ["""Subsample""", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
          ["""Subsample""", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
      ]
      _UpperCAmelCase         =					is_training
      _UpperCAmelCase         =					use_labels
      _UpperCAmelCase         =					num_labels
      _UpperCAmelCase         =					initializer_range


    def    lowerCAmelCase_     (    self		:     Dict    ):
      _UpperCAmelCase         =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

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


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      return LevitConfig(
          image_size=self.image_size      , num_channels=self.num_channels      , kernel_size=self.kernel_size      , stride=self.stride      , padding=self.padding      , patch_size=self.patch_size      , hidden_sizes=self.hidden_sizes      , num_attention_heads=self.num_attention_heads      , depths=self.depths      , key_dim=self.key_dim      , drop_path_rate=self.drop_path_rate      , mlp_ratio=self.mlp_ratio      , attention_ratio=self.attention_ratio      , initializer_range=self.initializer_range      , down_ops=self.down_ops      , )


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     List[str]      , __lowerCAmelCase		:     str    ):
      _UpperCAmelCase         =					LevitModel(config=__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase    )
      _UpperCAmelCase         =					(self.image_size, self.image_size)
      _UpperCAmelCase       , _UpperCAmelCase         =					image_size[0], image_size[1]
      for _ in range(4    ):
        _UpperCAmelCase         =					floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1    )
        _UpperCAmelCase         =					floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1    )
      self.parent.assertEqual(
          result.last_hidden_state.shape      , (self.batch_size, ceil(height / 4    ) * ceil(width / 4    ), self.hidden_sizes[-1])      , )


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     Any      , __lowerCAmelCase		:     Union[str, Any]      , __lowerCAmelCase		:     Union[str, Any]    ):
      _UpperCAmelCase         =					self.num_labels
      _UpperCAmelCase         =					LevitForImageClassification(__lowerCAmelCase    )
      model.to(__lowerCAmelCase    )
      model.eval()
      _UpperCAmelCase         =					model(__lowerCAmelCase      , labels=__lowerCAmelCase    )
      self.parent.assertEqual(result.logits.shape      , (self.batch_size, self.num_labels)    )


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase         =					self.prepare_config_and_inputs()
      _UpperCAmelCase       , _UpperCAmelCase       , _UpperCAmelCase         =					config_and_inputs
      _UpperCAmelCase         =					{"""pixel_values""": pixel_values}
      return config, inputs_dict



@require_torch
class 				a (						lowerCAmelCase_		,							lowerCAmelCase_		,							unittest.TestCase      ):
    _snake_case   :     Dict   							=							(
        (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
        if is_torch_available()
        else ()
    )
    _snake_case   :     Tuple   							=							(
        {
            'feature-extraction': LevitModel,
            'image-classification': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
        }
        if is_torch_available()
        else {}
    )

    _snake_case   :     Union[str, Any]   							=							False
    _snake_case   :     List[Any]   							=							False
    _snake_case   :     Any   							=							False
    _snake_case   :     List[str]   							=							False
    _snake_case   :     Tuple   							=							False


    def    lowerCAmelCase_     (    self		:     List[str]    ):
      _UpperCAmelCase         =					LevitModelTester(self    )
      _UpperCAmelCase         =					ConfigTester(self      , config_class=__lowerCAmelCase      , has_text_modality=__lowerCAmelCase      , hidden_size=37    )


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      self.create_and_test_config_common_properties()
      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    lowerCAmelCase_     (    self		:     int    ):
      return


    @unittest.skip(reason="""Levit does not use inputs_embeds"""    )
    def    lowerCAmelCase_     (    self		:     int    ):
      pass


    @unittest.skip(reason="""Levit does not support input and output embeddings"""    )
    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      pass


    @unittest.skip(reason="""Levit does not output attentions"""    )
    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      pass


    def    lowerCAmelCase_     (    self		:     List[Any]    ):
      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()

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

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


    def    lowerCAmelCase_     (    self		:     Dict    ):
      def check_hidden_states_output(__lowerCAmelCase		:     str      , __lowerCAmelCase		:     Optional[Any]      , __lowerCAmelCase		:     Tuple    ):
        _UpperCAmelCase         =					model_class(__lowerCAmelCase    )
        model.to(__lowerCAmelCase    )
        model.eval()

        with torch.no_grad():
          _UpperCAmelCase         =					model(**self._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase    )    )

        _UpperCAmelCase         =					outputs.hidden_states

        _UpperCAmelCase         =					len(self.model_tester.depths    ) + 1
        self.assertEqual(len(__lowerCAmelCase    )      , __lowerCAmelCase    )

        _UpperCAmelCase         =					(self.model_tester.image_size, self.model_tester.image_size)
        _UpperCAmelCase       , _UpperCAmelCase         =					image_size[0], image_size[1]
        for _ in range(4    ):
          _UpperCAmelCase         =					floor(
              (
                  (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                  / self.model_tester.stride
              )
              + 1    )
          _UpperCAmelCase         =					floor(
              (
                  (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                  / self.model_tester.stride
              )
              + 1    )
        # verify the first hidden states (first block)
        self.assertListEqual(
            list(hidden_states[0].shape[-2:]    )      , [
                height * width,
                self.model_tester.hidden_sizes[0],
            ]      , )

      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        _UpperCAmelCase         =					True
        check_hidden_states_output(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase    )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        _UpperCAmelCase         =					True

        check_hidden_states_output(__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase    )


    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""    )
    def    lowerCAmelCase_     (    self		:     Optional[Any]    ):
      pass


    def    lowerCAmelCase_     (    self		:     Optional[int]      , __lowerCAmelCase		:     Tuple      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     List[str]=False    ):
      _UpperCAmelCase         =					super()._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )

      if return_labels:
        if model_class.__name__ == "LevitForImageClassificationWithTeacher":
          del inputs_dict["labels"]

      return inputs_dict


    def    lowerCAmelCase_     (    self		:     Tuple    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__lowerCAmelCase    )


    def    lowerCAmelCase_     (    self		:     Optional[int]    ):
      if not self.model_tester.is_training:
        return

      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()
      _UpperCAmelCase         =					True

      for model_class in self.all_model_classes:
        # LevitForImageClassificationWithTeacher supports inference-only
        if (
            model_class in get_values(__lowerCAmelCase    )
            or model_class.__name__ == "LevitForImageClassificationWithTeacher"
        ):
          continue
        _UpperCAmelCase         =					model_class(__lowerCAmelCase    )
        model.to(__lowerCAmelCase    )
        model.train()
        _UpperCAmelCase         =					self._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )
        _UpperCAmelCase         =					model(**__lowerCAmelCase    ).loss
        loss.backward()


    def    lowerCAmelCase_     (    self		:     int    ):
      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()
      if not self.model_tester.is_training:
        return

      _UpperCAmelCase         =					False
      _UpperCAmelCase         =					True

      for model_class in self.all_model_classes:
        if model_class in get_values(__lowerCAmelCase    ) or not model_class.supports_gradient_checkpointing:
          continue
        # LevitForImageClassificationWithTeacher supports inference-only
        if model_class.__name__ == "LevitForImageClassificationWithTeacher":
          continue
        _UpperCAmelCase         =					model_class(__lowerCAmelCase    )
        model.gradient_checkpointing_enable()
        model.to(__lowerCAmelCase    )
        model.train()
        _UpperCAmelCase         =					self._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )
        _UpperCAmelCase         =					model(**__lowerCAmelCase    ).loss
        loss.backward()


    def    lowerCAmelCase_     (    self		:     str    ):
      _UpperCAmelCase       , _UpperCAmelCase         =					self.model_tester.prepare_config_and_inputs_for_common()

      _UpperCAmelCase         =					[
          {"""title""": """multi_label_classification""", """num_labels""": 2, """dtype""": torch.float},
          {"""title""": """single_label_classification""", """num_labels""": 1, """dtype""": torch.long},
          {"""title""": """regression""", """num_labels""": 1, """dtype""": torch.float},
      ]

      for model_class in self.all_model_classes:
        if (
            model_class
            not in [
                *get_values(__lowerCAmelCase    ),
            ]
            or model_class.__name__ == "LevitForImageClassificationWithTeacher"
        ):
          continue

        for problem_type in problem_types:
          with self.subTest(msg=f'''Testing {model_class} with {problem_type["title"]}'''    ):
            _UpperCAmelCase         =					problem_type["""title"""]
            _UpperCAmelCase         =					problem_type["""num_labels"""]

            _UpperCAmelCase         =					model_class(__lowerCAmelCase    )
            model.to(__lowerCAmelCase    )
            model.train()

            _UpperCAmelCase         =					self._prepare_for_class(__lowerCAmelCase      , __lowerCAmelCase      , return_labels=__lowerCAmelCase    )

            if problem_type["num_labels"] > 1:
              _UpperCAmelCase         =					inputs["""labels"""].unsqueeze(1    ).repeat(1      , problem_type["""num_labels"""]    )

            _UpperCAmelCase         =					inputs["""labels"""].to(problem_type["""dtype"""]    )

            # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
            # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
            # they have the same size." which is a symptom something in wrong for the regression problem.
            # See https://github.com/huggingface/transformers/issues/11780
            with warnings.catch_warnings(record=__lowerCAmelCase    ) as warning_list:
              _UpperCAmelCase         =					model(**__lowerCAmelCase    ).loss
            for w in warning_list:
              if "Using a target size that is different to the input size" in str(w.message    ):
                raise ValueError(
                    f'''Something is going wrong in the regression problem: intercepted {w.message}'''    )

            loss.backward()


    @slow
    def    lowerCAmelCase_     (    self		:     List[str]    ):
      for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        _UpperCAmelCase         =					LevitModel.from_pretrained(__lowerCAmelCase    )
        self.assertIsNotNone(__lowerCAmelCase    )






def 			__UpperCAmelCase    (	):







  """simple docstring"""
  _UpperCAmelCase         =					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""						)
  return image


@require_torch
@require_vision
class 				a (						unittest.TestCase      ):


    @cached_property
    def    lowerCAmelCase_     (    self		:     int    ):
      return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )


    @slow
    def    lowerCAmelCase_     (    self		:     int    ):
      _UpperCAmelCase         =					LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    ).to(
          __lowerCAmelCase    )

      _UpperCAmelCase         =					self.default_image_processor
      _UpperCAmelCase         =					prepare_img()
      _UpperCAmelCase         =					image_processor(images=__lowerCAmelCase      , return_tensors="""pt"""    ).to(__lowerCAmelCase    )

      # forward pass
      with torch.no_grad():
        _UpperCAmelCase         =					model(**__lowerCAmelCase    )

      # verify the logits
      _UpperCAmelCase         =					torch.Size((1, 1000)    )
      self.assertEqual(outputs.logits.shape      , __lowerCAmelCase    )

      _UpperCAmelCase         =					torch.tensor([1.0_448, -0.3_745, -1.8_317]    ).to(__lowerCAmelCase    )

      self.assertTrue(torch.allclose(outputs.logits[0, :3]      , __lowerCAmelCase      , atol=1e-4    )    )

"""simple docstring"""


from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers
