

from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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 tensorflow as tf

       from transformers import TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
       from PIL import Image

       from transformers import ConvNextImageProcessor







class 						__magic_name__      :



   def __init__(    self,   lowercase_,   lowercase_=13,   lowercase_=32,   lowercase_=3,   lowercase_=4,   lowercase_=[10, 20, 30, 40],   lowercase_=[2, 2, 3, 2],   lowercase_=True,   lowercase_=True,   lowercase_=37,   lowercase_="gelu",   lowercase_=10,   lowercase_=0.02,   lowercase_=3,   lowercase_=None,   )							->   Optional[int]:






     """simple docstring"""






     a__  =parent
     a__  =batch_size
     a__  =image_size
     a__  =num_channels
     a__  =num_stages
     a__  =hidden_sizes
     a__  =depths
     a__  =is_training
     a__  =use_labels
     a__  =intermediate_size
     a__  =hidden_act
     a__  =type_sequence_label_size
     a__  =initializer_range
     a__  =scope



   def 			_UpperCAmelCase			(    self	)							->   str:






     """simple docstring"""






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

     a__  =None
     if self.use_labels:
       a__  =ids_tensor([self.batch_size],   self.type_sequence_label_size	)

     a__  =self.get_config()

     return config, pixel_values, labels



   def 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     return ConvNextConfig(
         num_channels=self.num_channels,   hidden_sizes=self.hidden_sizes,   depths=self.depths,   num_stages=self.num_stages,   hidden_act=self.hidden_act,   is_decoder=lowercase_,   initializer_range=self.initializer_range,   )



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_,   lowercase_	)							->   Tuple:






     """simple docstring"""






     a__  =TFConvNextModel(config=lowercase_	)
     a__  =model(lowercase_,   training=lowercase_	)
     # expected last hidden states: B, C, H // 32, W // 32
     self.parent.assertEqual(
         result.last_hidden_state.shape,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32),   )



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_,   lowercase_	)							->   List[Any]:






     """simple docstring"""






     a__  =self.type_sequence_label_size
     a__  =TFConvNextForImageClassification(lowercase_	)
     a__  =model(lowercase_,   labels=lowercase_,   training=lowercase_	)
     self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.type_sequence_label_size)	)



   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






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







@require_tf
class 						__magic_name__      (							lowerCamelCase__		,	lowerCamelCase__		,	unittest.TestCase							):
   lowerCamelCase__			:     List[Any]    =       (TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
   lowerCamelCase__			:     str    =       (
       {'feature-extraction': TFConvNextModel, 'image-classification': TFConvNextForImageClassification}
       if is_tf_available()
       else {}
   )

   lowerCamelCase__			:     Tuple    =       False
   lowerCamelCase__			:     str    =       False
   lowerCamelCase__			:     Union[str, Any]    =       False
   lowerCamelCase__			:     Dict    =       False
   lowerCamelCase__			:     Optional[Any]    =       False



   def 			_UpperCAmelCase			(    self	)							->   Optional[int]:






     """simple docstring"""






     a__  =TFConvNextModelTester(self	)
     a__  =ConfigTester(
         self,   config_class=lowercase_,   has_text_modality=lowercase_,   hidden_size=37,   )



   @unittest.skip(reason='''ConvNext does not use inputs_embeds'''	)
   def 			_UpperCAmelCase			(    self	)							->   Dict:






     """simple docstring"""






     pass



   @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 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     super().test_keras_fit()



   @unittest.skip(reason='''ConvNext does not support input and output embeddings'''	)
   def 			_UpperCAmelCase			(    self	)							->   Dict:






     """simple docstring"""






     pass



   def 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     a__,				a__  =self.model_tester.prepare_config_and_inputs_for_common()

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

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



   def 			_UpperCAmelCase			(    self	)							->   int:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_	)



   @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 			_UpperCAmelCase			(    self	)							->   str:






     """simple docstring"""






     super().test_dataset_conversion()



   def 			_UpperCAmelCase			(    self	)							->   Union[str, Any]:






     """simple docstring"""






     def check_hidden_states_output(lowercase_,   lowercase_,   lowercase_	):
       a__  =model_class(lowercase_	)

       a__  =model(**self._prepare_for_class(lowercase_,   lowercase_	)	)
       a__  =outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

       a__  =self.model_tester.num_stages
       self.assertEqual(len(lowercase_	),   expected_num_stages + 1	)

       # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
       self.assertListEqual(
           list(hidden_states[0].shape[-2:]	),   [self.model_tester.image_size // 4, self.model_tester.image_size // 4],   )

     a__,				a__  =self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       a__  =True
       check_hidden_states_output(lowercase_,   lowercase_,   lowercase_	)

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

       check_hidden_states_output(lowercase_,   lowercase_,   lowercase_	)



   def 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     a__,				a__  =self.model_tester.prepare_config_and_inputs_for_common()

     def check_equivalence(lowercase_,   lowercase_,   lowercase_,   lowercase_={}	):
       a__  =model(lowercase_,   return_dict=lowercase_,   **lowercase_	)
       a__  =model(lowercase_,   return_dict=lowercase_,   **lowercase_	).to_tuple()

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

         recursive_check(lowercase_,   lowercase_	)

     for model_class in self.all_model_classes:
       a__  =model_class(lowercase_	)

       a__  =self._prepare_for_class(lowercase_,   lowercase_	)
       a__  =self._prepare_for_class(lowercase_,   lowercase_	)
       check_equivalence(lowercase_,   lowercase_,   lowercase_	)

       a__  =self._prepare_for_class(lowercase_,   lowercase_,   return_labels=lowercase_	)
       a__  =self._prepare_for_class(lowercase_,   lowercase_,   return_labels=lowercase_	)
       check_equivalence(lowercase_,   lowercase_,   lowercase_	)

       a__  =self._prepare_for_class(lowercase_,   lowercase_	)
       a__  =self._prepare_for_class(lowercase_,   lowercase_	)
       check_equivalence(lowercase_,   lowercase_,   lowercase_,   {'''output_hidden_states''': True}	)

       a__  =self._prepare_for_class(lowercase_,   lowercase_,   return_labels=lowercase_	)
       a__  =self._prepare_for_class(lowercase_,   lowercase_,   return_labels=lowercase_	)
       check_equivalence(lowercase_,   lowercase_,   lowercase_,   {'''output_hidden_states''': True}	)



   def 			_UpperCAmelCase			(    self	)							->   int:






     """simple docstring"""






     a__  =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_	)



   @slow
   def 			_UpperCAmelCase			(    self	)							->   str:






     """simple docstring"""






     a__  =TFConvNextModel.from_pretrained('''facebook/convnext-tiny-224'''	)
     self.assertIsNotNone(lowercase_	)





def       UpperCAmelCase__    (      ):





  '''simple docstring'''







  a__  =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
  return image







@require_tf
@require_vision
class 						__magic_name__      (							unittest.TestCase							):



   @cached_property
   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






     return ConvNextImageProcessor.from_pretrained('''facebook/convnext-tiny-224'''	) if is_vision_available() else None



   @slow
   def 			_UpperCAmelCase			(    self	)							->   int:






     """simple docstring"""






     a__  =TFConvNextForImageClassification.from_pretrained('''facebook/convnext-tiny-224'''	)

     a__  =self.default_image_processor
     a__  =prepare_img()
     a__  =image_processor(images=lowercase_,   return_tensors='''tf'''	)

     # forward pass
     a__  =model(**lowercase_	)

     # verify the logits
     a__  =tf.TensorShape((1, 1000)	)
     self.assertEqual(outputs.logits.shape,   lowercase_	)

     a__  =tf.constant([-0.0260, -0.4739, 0.1911]	)

     tf.debugging.assert_near(outputs.logits[0, :3],   lowercase_,   atol=1E-4	)



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


lowerCamelCase									= get_logger(__name__)




def       UpperCAmelCase__    (      _A				:			str     ,				_A				:			List[dict]     ,				_A				:			list     ,				_A				:			list     ,				_A				:			int     ,				_A				:			GradScaler = None     ,				):





  '''simple docstring'''







  for i, state in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Model weights saved in {output_model_file}"""							)
  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =opt.state_dict()
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Optimizer state saved in {output_optimizer_file}"""							)
  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =scheduler.state_dict()
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Scheduler state saved in {output_scheduler_file}"""							)
  # GradScaler state
  if scaler is not None:
    a__  =scaler.state_dict()
    a__  =os.path.join(_A     ,				_A							)
    torch.save(_A     ,				_A							)
    logger.info(F"""Gradient scaler state saved in {output_scaler_file}"""							)
  # Random number generator states
  a__  ={}
  a__  =F"""{RNG_STATE_NAME}_{process_index}.pkl"""
  a__  =random.getstate()
  a__  =np.random.get_state()
  a__  =torch.get_rng_state()
  if is_xpu_available():
    a__  =torch.xpu.get_rng_state_all()
  else:
    a__  =torch.cuda.get_rng_state_all()
  if is_tpu_available():
    a__  =xm.get_rng_state()
  a__  =os.path.join(_A     ,				_A							)
  torch.save(_A     ,				_A							)
  logger.info(F"""Random states saved in {output_states_file}"""							)
  return output_dir




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			int     ,				_A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			List[str]=None     ,				_A				:			Dict=None     ,				**_A				:			Tuple     ,				):





  '''simple docstring'''







  if map_location not in [None, "cpu", "on_device"]:
    raise TypeError(
        '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
  if map_location is None:
    a__  ='''cpu'''
  elif map_location == "on_device":
    a__  =PartialState().device
  # Model states
  for i, model in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    models[i].load_state_dict(torch.load(_A     ,				map_location=_A							)     ,				**_A							)
  logger.info('''All model weights loaded successfully'''							)

  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    a__  =torch.load(_A     ,				map_location=_A							)
    optimizers[i].load_state_dict(_A							)
  logger.info('''All optimizer states loaded successfully'''							)

  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    scheduler.load_state_dict(torch.load(_A							)							)
  logger.info('''All scheduler states loaded successfully'''							)

  # GradScaler state
  if scaler is not None:
    a__  =os.path.join(_A     ,				_A							)
    scaler.load_state_dict(torch.load(_A							)							)
    logger.info('''GradScaler state loaded successfully'''							)

  # Random states
  try:
    a__  =torch.load(os.path.join(_A     ,				F"""{RNG_STATE_NAME}_{process_index}.pkl"""							)							)
    random.setstate(states['''random_state''']							)
    np.random.set_state(states['''numpy_random_seed''']							)
    torch.set_rng_state(states['''torch_manual_seed''']							)
    if is_xpu_available():
      torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
    else:
      torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
    if is_tpu_available():
      xm.set_rng_state(states['''xm_seed''']							)
    logger.info('''All random states loaded successfully'''							)
  except Exception:
    logger.info('''Could not load random states'''							)




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =Path(_A							) / F"""custom_checkpoint_{index}.pkl"""
  logger.info(F"""Saving the state of {get_pretty_name(_A							)} to {save_location}"""							)
  torch.save(obj.state_dict()     ,				_A							)




def       UpperCAmelCase__    (      _A				:			int     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =F"""{path}/custom_checkpoint_{index}.pkl"""
  logger.info(F"""Loading the state of {get_pretty_name(_A							)} from {load_location}"""							)
  obj.load_state_dict(torch.load(_A     ,				map_location='''cpu'''							)							)
