





import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
       import jax
       import jax.numpy as jnp

       from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor





class A_   (     unittest.TestCase  ):




   def __init__(self      :List[Any]       ,  _UpperCamelCase      :Union[str, Any]       ,  _UpperCamelCase      :str=3       ,  _UpperCamelCase      :List[Any]=32       ,  _UpperCamelCase      :Tuple=3       ,  _UpperCamelCase      :Union[str, Any]=10       ,  _UpperCamelCase      :Dict=[10, 20, 30, 40]       ,  _UpperCamelCase      :Dict=[1, 1, 2, 1]       ,  _UpperCamelCase      :Union[str, Any]=True       ,  _UpperCamelCase      :int=True       ,  _UpperCamelCase      :str="relu"       ,  _UpperCamelCase      :List[str]=3       ,  _UpperCamelCase      :int=None       ,  )->   List[str]:
         __A			  =							parent
         __A			  =							batch_size
         __A			  =							image_size
         __A			  =							num_channels
         __A			  =							embeddings_size
         __A			  =							hidden_sizes
         __A			  =							depths
         __A			  =							is_training
         __A			  =							use_labels
         __A			  =							hidden_act
         __A			  =							num_labels
         __A			  =							scope
         __A			  =							len(_UpperCamelCase					)




   def _lowerCAmelCase (self      :Any					)->   str:
         __A			  =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

         __A			  =							self.get_config()

         return config, pixel_values




   def _lowerCAmelCase (self      :int					)->   List[str]:
         return ResNetConfig(
             num_channels=self.num_channels       ,  embeddings_size=self.embeddings_size       ,  hidden_sizes=self.hidden_sizes       ,  depths=self.depths       ,  hidden_act=self.hidden_act       ,  num_labels=self.num_labels       ,  image_size=self.image_size       ,  )




   def _lowerCAmelCase (self      :Optional[int]       ,  _UpperCamelCase      :List[Any]       ,  _UpperCamelCase      :Any					)->   Optional[Any]:
         __A			  =							FlaxResNetModel(config=_UpperCamelCase					)
         __A			  =							model(_UpperCamelCase					)

         # Output shape (b, c, h, w)
         self.parent.assertEqual(
             result.last_hidden_state.shape       ,  (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)       ,  )




   def _lowerCAmelCase (self      :List[str]       ,  _UpperCamelCase      :Any       ,  _UpperCamelCase      :Union[str, Any]					)->   List[str]:
         __A			  =							self.num_labels
         __A			  =							FlaxResNetForImageClassification(config=_UpperCamelCase					)
         __A			  =							model(_UpperCamelCase					)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)					)






   def _lowerCAmelCase (self      :Union[str, Any]					)->   Dict:
         __A			  =							self.prepare_config_and_inputs()
         __A  ,		__A			  =							config_and_inputs
         __A			  =							{'''pixel_values''': pixel_values}
         return config, inputs_dict





@require_flax
class A_   (     _lowerCamelCase					, unittest.TestCase  ):
   lowerCAmelCase__  	=					(FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

   lowerCAmelCase__  	=					False
   lowerCAmelCase__  	=					False
   lowerCAmelCase__  	=					False




   def _lowerCAmelCase (self      :Any					)->   None:
         __A			  =							FlaxResNetModelTester(self					)
         __A			  =							ConfigTester(self       ,  config_class=_UpperCamelCase       ,  has_text_modality=_UpperCamelCase					)




   def _lowerCAmelCase (self      :Optional[Any]					)->   int:
         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      :Any					)->   Dict:
         return




   def _lowerCAmelCase (self      :Optional[Any]					)->   List[Any]:
         __A			  =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_UpperCamelCase					)




   def _lowerCAmelCase (self      :Optional[Any]					)->   Any:
         __A			  =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase					)




   @unittest.skip(reason='''ResNet does not use inputs_embeds'''					)
   def _lowerCAmelCase (self      :Dict					)->   Tuple:
         pass




   @unittest.skip(reason='''ResNet does not support input and output embeddings'''					)
   def _lowerCAmelCase (self      :Union[str, Any]					)->   Dict:
         pass




   def _lowerCAmelCase (self      :str					)->   Dict:
         __A  ,		__A			  =							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
               __A			  =							model_class(_UpperCamelCase					)
               __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]       ,  _UpperCamelCase					)




   def _lowerCAmelCase (self      :Optional[Any]					)->   Any:
         def check_hidden_states_output(_UpperCamelCase      :Any       ,  _UpperCamelCase      :str       ,  _UpperCamelCase      :int					):
               __A			  =							model_class(_UpperCamelCase					)

               __A			  =							model(**self._prepare_for_class(_UpperCamelCase       ,  _UpperCamelCase					)					)

               __A			  =							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               __A			  =							self.model_tester.num_stages
               self.assertEqual(len(_UpperCamelCase					)       ,  expected_num_stages + 1					)




   @unittest.skip(reason='''ResNet does not use feedforward chunking'''					)
   def _lowerCAmelCase (self      :str					)->   Optional[Any]:
         pass






   def _lowerCAmelCase (self      :Tuple					)->   Tuple:
         __A  ,		__A			  =							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__					):
                     __A			  =							self._prepare_for_class(_UpperCamelCase       ,  _UpperCamelCase					)
                     __A			  =							model_class(_UpperCamelCase					)

                     @jax.jit
                     def model_jitted(_UpperCamelCase      :Dict       ,  **_UpperCamelCase      :Any					):
                           return model(pixel_values=_UpperCamelCase       ,  **_UpperCamelCase					)

                     with self.subTest('''JIT Enabled'''					):
                           __A			  =							model_jitted(**_UpperCamelCase					).to_tuple()

                     with self.subTest('''JIT Disabled'''					):
                           with jax.disable_jit():
                                 __A			  =							model_jitted(**_UpperCamelCase					).to_tuple()

                     self.assertEqual(len(_UpperCamelCase					)       ,  len(_UpperCamelCase					)					)
                     for jitted_output, output in zip(_UpperCamelCase       ,  _UpperCamelCase					):
                           self.assertEqual(jitted_output.shape       ,  output.shape					)





def _a							(      )      ->		Any:







      '''simple docstring'''






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







@require_flax
class A_   (     unittest.TestCase  ):




   @cached_property
   def _lowerCAmelCase (self      :Union[str, Any]					)->   List[str]:
         return AutoImageProcessor.from_pretrained('''microsoft/resnet-50'''					) if is_vision_available() else None






   @slow
   def _lowerCAmelCase (self      :Union[str, Any]					)->   Any:
         __A			  =							FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''					)

         __A			  =							self.default_image_processor
         __A			  =							prepare_img()
         __A			  =							image_processor(images=_UpperCamelCase       ,  return_tensors='''np'''					)

         __A			  =							model(**_UpperCamelCase					)

         # verify the logits
         __A			  =							(1, 1000)
         self.assertEqual(outputs.logits.shape       ,  _UpperCamelCase					)

         __A			  =							jnp.array([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]					)

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












import torch

from diffusers import LMSDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest





class A_   (     _lowerCamelCase  ):
   lowerCAmelCase__  	=					(LMSDiscreteScheduler,)
   lowerCAmelCase__  	=					10




   def _lowerCAmelCase (self      :List[str]       ,  **_UpperCamelCase      :Optional[int]					)->   List[str]:
         __A			  =							{
             '''num_train_timesteps''': 1100,
             '''beta_start''': 0.0_0_0_1,
             '''beta_end''': 0.0_2,
             '''beta_schedule''': '''linear''',
         }

         config.update(**_UpperCamelCase					)
         return config




   def _lowerCAmelCase (self      :List[str]					)->   Tuple:
         for timesteps in [10, 50, 100, 1000]:
               self.check_over_configs(num_train_timesteps=_UpperCamelCase					)




   def _lowerCAmelCase (self      :Dict					)->   Optional[int]:
         for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]       ,  [0.0_0_0_2, 0.0_0_2, 0.0_2]					):
               self.check_over_configs(beta_start=_UpperCamelCase       ,  beta_end=_UpperCamelCase					)




   def _lowerCAmelCase (self      :Optional[int]					)->   Any:
         for schedule in ["linear", "scaled_linear"]:
               self.check_over_configs(beta_schedule=_UpperCamelCase					)




   def _lowerCAmelCase (self      :List[str]					)->   str:
         for prediction_type in ["epsilon", "v_prediction"]:
               self.check_over_configs(prediction_type=_UpperCamelCase					)




   def _lowerCAmelCase (self      :str					)->   List[str]:
         for t in [0, 500, 800]:
               self.check_over_forward(time_step=_UpperCamelCase					)




   def _lowerCAmelCase (self      :Optional[Any]					)->   Tuple:
         __A			  =							self.scheduler_classes[0]
         __A			  =							self.get_scheduler_config()
         __A			  =							scheduler_class(**_UpperCamelCase					)

         scheduler.set_timesteps(self.num_inference_steps					)

         __A			  =							self.dummy_model()
         __A			  =							self.dummy_sample_deter * scheduler.init_noise_sigma

         for i, t in enumerate(scheduler.timesteps					):
               __A			  =							scheduler.scale_model_input(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							model(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							scheduler.step(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase					)
               __A			  =							output.prev_sample

         __A			  =							torch.sum(torch.abs(_UpperCamelCase					)					)
         __A			  =							torch.mean(torch.abs(_UpperCamelCase					)					)

         assert abs(result_sum.item() - 1_0_0_6.3_8_8					) < 1e-2
         assert abs(result_mean.item() - 1.3_1					) < 1e-3




   def _lowerCAmelCase (self      :Optional[Any]					)->   Optional[int]:
         __A			  =							self.scheduler_classes[0]
         __A			  =							self.get_scheduler_config(prediction_type='''v_prediction'''					)
         __A			  =							scheduler_class(**_UpperCamelCase					)

         scheduler.set_timesteps(self.num_inference_steps					)

         __A			  =							self.dummy_model()
         __A			  =							self.dummy_sample_deter * scheduler.init_noise_sigma

         for i, t in enumerate(scheduler.timesteps					):
               __A			  =							scheduler.scale_model_input(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							model(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							scheduler.step(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase					)
               __A			  =							output.prev_sample

         __A			  =							torch.sum(torch.abs(_UpperCamelCase					)					)
         __A			  =							torch.mean(torch.abs(_UpperCamelCase					)					)

         assert abs(result_sum.item() - 0.0_0_1_7					) < 1e-2
         assert abs(result_mean.item() - 2.2_676e-06					) < 1e-3




   def _lowerCAmelCase (self      :List[Any]					)->   List[Any]:
         __A			  =							self.scheduler_classes[0]
         __A			  =							self.get_scheduler_config()
         __A			  =							scheduler_class(**_UpperCamelCase					)

         scheduler.set_timesteps(self.num_inference_steps       ,  device=_UpperCamelCase					)

         __A			  =							self.dummy_model()
         __A			  =							self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
         __A			  =							sample.to(_UpperCamelCase					)

         for i, t in enumerate(scheduler.timesteps					):
               __A			  =							scheduler.scale_model_input(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							model(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							scheduler.step(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase					)
               __A			  =							output.prev_sample

         __A			  =							torch.sum(torch.abs(_UpperCamelCase					)					)
         __A			  =							torch.mean(torch.abs(_UpperCamelCase					)					)

         assert abs(result_sum.item() - 1_0_0_6.3_8_8					) < 1e-2
         assert abs(result_mean.item() - 1.3_1					) < 1e-3






   def _lowerCAmelCase (self      :List[str]					)->   int:
         __A			  =							self.scheduler_classes[0]
         __A			  =							self.get_scheduler_config()
         __A			  =							scheduler_class(**_UpperCamelCase       ,  use_karras_sigmas=_UpperCamelCase					)

         scheduler.set_timesteps(self.num_inference_steps       ,  device=_UpperCamelCase					)

         __A			  =							self.dummy_model()
         __A			  =							self.dummy_sample_deter.to(_UpperCamelCase					) * scheduler.init_noise_sigma
         __A			  =							sample.to(_UpperCamelCase					)

         for t in scheduler.timesteps:
               __A			  =							scheduler.scale_model_input(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							model(_UpperCamelCase       ,  _UpperCamelCase					)

               __A			  =							scheduler.step(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase					)
               __A			  =							output.prev_sample

         __A			  =							torch.sum(torch.abs(_UpperCamelCase					)					)
         __A			  =							torch.mean(torch.abs(_UpperCamelCase					)					)

         assert abs(result_sum.item() - 3_8_1_2.9_9_2_7					) < 2e-2
         assert abs(result_mean.item() - 4.9_6_4_8					) < 1e-3





