
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							):



       """simple docstring"""






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






       def 	__lowerCAmelCase     (		self   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Any				      = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.get_config()

           return config, pixel_values






       def 	__lowerCAmelCase     (		self   )					->Any:
           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				,       _lowerCamelCase				,       _lowerCamelCase   )					->str:
           SCREAMING_SNAKE_CASE				:				Any				      = FlaxResNetModel(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				str				      = model(_lowerCamelCase   )

           # 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				,       _lowerCamelCase				,       _lowerCamelCase   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.num_labels
           SCREAMING_SNAKE_CASE				:				str				      = FlaxResNetForImageClassification(config=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = model(_lowerCamelCase   )
           self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)   )






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






@require_flax
class     a_       (			a__				,    unittest.TestCase							):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE       :					Union[str, Any]						=  (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

       __SCREAMING_SNAKE_CASE       :					Optional[int]						=  False
       __SCREAMING_SNAKE_CASE       :					List[Any]						=  False
       __SCREAMING_SNAKE_CASE       :					Tuple						=  False






       def 	__lowerCAmelCase     (		self   )					->None:
           SCREAMING_SNAKE_CASE				:				List[str]				      = FlaxResNetModelTester(self   )
           SCREAMING_SNAKE_CASE				:				str				      = ConfigTester(self				,       config_class=_lowerCamelCase				,       has_text_modality=_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->str:
           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   )					->Union[str, Any]:
           return






       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				str				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Dict:
           SCREAMING_SNAKE_CASE				:				Tuple				      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase   )






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






       @unittest.skip(reason='''ResNet does not support input and output embeddings'''   )
       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           pass






       def 	__lowerCAmelCase     (		self   )					->int:
           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Optional[int]				      = self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               SCREAMING_SNAKE_CASE				:				List[str]				      = model_class(_lowerCamelCase   )
               SCREAMING_SNAKE_CASE				:				List[Any]				      = inspect.signature(model.__call__   )
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               SCREAMING_SNAKE_CASE				:				int				      = [*signature.parameters.keys()]

               SCREAMING_SNAKE_CASE				:				List[str]				      = ['''pixel_values''']
               self.assertListEqual(arg_names[:1]				,       _lowerCamelCase   )






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           def check_hidden_states_output(_lowerCamelCase				,       _lowerCamelCase				,       _lowerCamelCase   ):
               SCREAMING_SNAKE_CASE				:				int				      = model_class(_lowerCamelCase   )

               SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = model(**self._prepare_for_class(_lowerCamelCase				,       _lowerCamelCase   )   )

               SCREAMING_SNAKE_CASE				:				Dict				      = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               SCREAMING_SNAKE_CASE				:				Optional[int]				      = self.model_tester.num_stages
               self.assertEqual(len(_lowerCamelCase   )				,       expected_num_stages + 1   )






       @unittest.skip(reason='''ResNet does not use feedforward chunking'''   )
       def 	__lowerCAmelCase     (		self   )					->Tuple:
           pass






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__   ):
                   SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self._prepare_for_class(_lowerCamelCase				,       _lowerCamelCase   )
                   SCREAMING_SNAKE_CASE				:				Optional[Any]				      = model_class(_lowerCamelCase   )

                   @jax.jit
                   def model_jitted(_lowerCamelCase				,       **_lowerCamelCase   ):
                       return model(pixel_values=_lowerCamelCase				,       **_lowerCamelCase   )

                   with self.subTest('''JIT Enabled'''   ):
                       SCREAMING_SNAKE_CASE				:				Optional[int]				      = model_jitted(**_lowerCamelCase   ).to_tuple()

                   with self.subTest('''JIT Disabled'''   ):
                       with jax.disable_jit():
                           SCREAMING_SNAKE_CASE				:				str				      = model_jitted(**_lowerCamelCase   ).to_tuple()

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







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[int]				      = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
    return image






@require_flax
class     a_       (			unittest.TestCase							):



       """simple docstring"""






       @cached_property
       def 	__lowerCAmelCase     (		self   )					->Optional[Any]:
           return AutoImageProcessor.from_pretrained('''microsoft/resnet-50'''   ) if is_vision_available() else None






       @slow
       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				List[str]				      = FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''   )

           SCREAMING_SNAKE_CASE				:				Dict				      = self.default_image_processor
           SCREAMING_SNAKE_CASE				:				List[Any]				      = prepare_img()
           SCREAMING_SNAKE_CASE				:				List[str]				      = image_processor(images=_lowerCamelCase				,       return_tensors='''np'''   )

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = model(**_lowerCamelCase   )

           # verify the logits
           SCREAMING_SNAKE_CASE				:				Any				      = (1, 1000)
           self.assertEqual(outputs.logits.shape				,       _lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				List[Any]				      = 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]				,       _lowerCamelCase				,       atol=1e-4   )   )


import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


a__				:      Tuple    		=     [nn.LayerNorm]

a__				:      Tuple    		=     logging.get_logger(__name__)

a__				:      int    		=     version.parse(version.parse(torch.__version__).base_version)

a__				:      Dict    		=     parsed_torch_version_base >= version.parse('''2.0''')
a__				:      Any    		=     parsed_torch_version_base >= version.parse('''1.12''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base >= version.parse('''1.11''')
a__				:      Union[str, Any]    		=     parsed_torch_version_base < version.parse('''1.11''')







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__ ,						self   ):



    """simple docstring"""



    from torch import _softmax_backward_data

    if is_torch_less_than_1_11:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self   )
    else:
        return _softmax_backward_data(a__ ,						a__ ,						parent.dim ,						self.dtype   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 0   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            SCREAMING_SNAKE_CASE				:				int				      = layer.bias.clone().detach()
        else:
            SCREAMING_SNAKE_CASE				:				List[Any]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				List[Any]				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				List[Any]				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				str				      = nn.Linear(new_size[1] ,						new_size[0] ,						bias=layer.bias is not None   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = True
    if layer.bias is not None:
        SCREAMING_SNAKE_CASE				:				Dict				      = False
        new_layer.bias.copy_(b.contiguous()   )
        SCREAMING_SNAKE_CASE				:				Tuple				      = True
    return new_layer






class     a_       (			nn.Module							):



       """simple docstring"""






       def __init__(		self				,       _lowerCamelCase				,       _lowerCamelCase   )					->int:
           super().__init__()
           SCREAMING_SNAKE_CASE				:				List[str]				      = nf
           SCREAMING_SNAKE_CASE				:				List[Any]				      = nn.Parameter(torch.empty(_lowerCamelCase				,       _lowerCamelCase   )   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = nn.Parameter(torch.zeros(_lowerCamelCase   )   )
           nn.init.normal_(self.weight				,       std=0.0_2   )






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase   )					->Tuple:
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = x.size()[:-1] + (self.nf,)
           SCREAMING_SNAKE_CASE				:				int				      = torch.addmm(self.bias				,       x.view(-1				,       x.size(-1   )   )				,       self.weight   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = x.view(_lowerCamelCase   )
           return x







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = 1   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				str				      = index.to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Any				      = layer.weight.index_select(a__ ,						a__   ).clone().detach()
    if dim == 0:
        SCREAMING_SNAKE_CASE				:				Tuple				      = layer.bias.clone().detach()
    else:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = layer.bias[index].clone().detach()
    SCREAMING_SNAKE_CASE				:				int				      = list(layer.weight.size()   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = len(a__   )
    SCREAMING_SNAKE_CASE				:				Any				      = ConvaD(new_size[1] ,						new_size[0]   ).to(layer.weight.device   )
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = False
    new_layer.weight.copy_(W.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    SCREAMING_SNAKE_CASE				:				Any				      = False
    new_layer.bias.copy_(b.contiguous()   )
    SCREAMING_SNAKE_CASE				:				int				      = True
    return new_layer







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ = None   ):



    """simple docstring"""



    if isinstance(a__ ,						nn.Linear   ):
        return prune_linear_layer(a__ ,						a__ ,						dim=0 if dim is None else dim   )
    elif isinstance(a__ ,						a__   ):
        return prune_convad_layer(a__ ,						a__ ,						dim=1 if dim is None else dim   )
    else:
        raise ValueError(F"""Can't prune layer of class {layer.__class__}"""   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						*a__   ):



    """simple docstring"""



    assert len(a__   ) > 0, F"""{input_tensors} has to be a tuple/list of tensors"""

    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    SCREAMING_SNAKE_CASE				:				str				      = len(inspect.signature(a__   ).parameters   )
    if num_args_in_forward_chunk_fn != len(a__   ):
        raise ValueError(
            F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(a__   )} input """
            '''tensors are given'''   )

    if chunk_size > 0:
        SCREAMING_SNAKE_CASE				:				Dict				      = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    F"""All input tenors have to be of the same shape: {tensor_shape}, """
                    F"""found shape {input_tensor.shape[chunk_dim]}"""   )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                F"""The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk """
                F"""size {chunk_size}"""   )

        SCREAMING_SNAKE_CASE				:				Any				      = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = tuple(input_tensor.chunk(a__ ,						dim=a__   ) for input_tensor in input_tensors   )
        # apply forward fn to every tuple
        SCREAMING_SNAKE_CASE				:				Tuple				      = tuple(forward_fn(*a__   ) for input_tensors_chunk in zip(*a__   )   )
        # concatenate output at same dimension
        return torch.cat(a__ ,						dim=a__   )

    return forward_fn(*a__   )







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[int]				      = torch.ones(a__ ,						a__   )
    SCREAMING_SNAKE_CASE				:				str				      = set(a__   ) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = head - sum(1 if h < head else 0 for h in already_pruned_heads   )
        SCREAMING_SNAKE_CASE				:				str				      = 0
    SCREAMING_SNAKE_CASE				:				List[str]				      = mask.view(-1   ).contiguous().eq(1   )
    SCREAMING_SNAKE_CASE				:				torch.LongTensor				      = torch.arange(len(a__   )   )[mask].long()
    return heads, index







def 				UpperCAmelCase_(    *a__ ,						a__ = None   ):



    """simple docstring"""



    return torch.meshgrid(*a__ ,						indexing=a__   )







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    return tensor.device, storage_ptr(a__   ), storage_size(a__   )
