





"""simple docstring"""




import inspect
import unittest

from transformers import EfficientNetConfig
from transformers.testing_utils import is_pipeline_test, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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 EfficientNetForImageClassification, EfficientNetModel
   from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor

class 					_UpperCAmelCase    :



      '''simple docstring'''

      def __init__(	self   , snake_case_   , snake_case_=1_3   , snake_case_=3_2   , snake_case_=3   , snake_case_=[3, 3, 5]   , snake_case_=[3_2, 1_6, 2_4]   , snake_case_=[1_6, 2_4, 4_0]   , snake_case_=[1, 1, 2]   , snake_case_=[1, 1, 2]   , snake_case_=[1, 6, 6]   , snake_case_=True   , snake_case_=True   , snake_case_=3_7   , snake_case_="gelu"   , snake_case_=1_0   , ):

        """simple docstring"""

        A_    :			int									=    parent
        A_    :			Dict									=    batch_size
        A_    :			Union[str, Any]									=    image_size
        A_    :			Optional[Any]									=    num_channels
        A_    :			Optional[Any]									=    kernel_sizes
        A_    :			List[Any]									=    in_channels
        A_    :			Tuple									=    out_channels
        A_    :			List[Any]									=    strides
        A_    :			Optional[Any]									=    num_block_repeats
        A_    :			Any									=    expand_ratios
        A_    :			Optional[int]									=    is_training
        A_    :			Tuple									=    hidden_act
        A_    :			Union[str, Any]									=    num_labels
        A_    :			str									=    use_labels

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Union[str, Any]									=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

        A_    :			Optional[int]									=    None
        if self.use_labels:
          A_    :			int									=    ids_tensor([self.batch_size]   , self.num_labels    )

        A_    :			List[Any]									=    self.get_config()
        return config, pixel_values, labels

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        return EfficientNetConfig(
            num_channels=self.num_channels   , kernel_sizes=self.kernel_sizes   , in_channels=self.in_channels   , out_channels=self.out_channels   , strides=self.strides   , num_block_repeats=self.num_block_repeats   , expand_ratios=self.expand_ratios   , hidden_act=self.hidden_act   , num_labels=self.num_labels   , )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			List[str]									=    EfficientNetModel(config=snake_case_    )
        model.to(snake_case_    )
        model.eval()
        A_    :			Dict									=    model(snake_case_    )
        # expected last hidden states: B, C, H // 4, W // 4
        self.parent.assertEqual(
            result.last_hidden_state.shape   , (self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)   , )

      def    lowerCamelCase_     (	self   , snake_case_   , snake_case_   , snake_case_    ):

        """simple docstring"""

        A_    :			int									=    EfficientNetForImageClassification(snake_case_    )
        model.to(snake_case_    )
        model.eval()
        A_    :			Optional[int]									=    model(snake_case_   , labels=snake_case_    )
        self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_labels)    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			int									=    self.prepare_config_and_inputs()
        A_		,			A_		,			A_    :			List[str]									=    config_and_inputs
        A_    :			Tuple									=    {'pixel_values': pixel_values}
        return config, inputs_dict


@require_torch
class 					_UpperCAmelCase    (  UpperCAmelCase__     ,  UpperCAmelCase__     ,  unittest.TestCase						):



      '''simple docstring'''
      lowercase_ :					Union[str, Any]             =  (EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
      lowercase_ :					List[str]             =  (
          {"""feature-extraction""": EfficientNetModel, """image-classification""": EfficientNetForImageClassification}
          if is_torch_available()
          else {}
      )

      lowercase_ :					Optional[Any]             =  False
      lowercase_ :					Optional[int]             =  False
      lowercase_ :					Optional[int]             =  False
      lowercase_ :					Tuple             =  False
      lowercase_ :					Optional[int]             =  False

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[Any]									=    EfficientNetModelTester(self    )
        A_    :			List[Any]									=    ConfigTester(
            self   , config_class=snake_case_   , has_text_modality=snake_case_   , hidden_size=3_7    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

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

        """simple docstring"""

        return

      @unittest.skip(reason='EfficientNet does not use inputs_embeds'    )
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        pass

      @unittest.skip(reason='EfficientNet does not support input and output embeddings'    )
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        pass

      @unittest.skip(reason='EfficientNet does not use feedforward chunking'    )
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        pass

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_		,			A_    :			str									=    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          A_    :			Tuple									=    model_class(snake_case_    )
          A_    :			Optional[Any]									=    inspect.signature(model.forward    )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          A_    :			Dict									=    [*signature.parameters.keys()]

          A_    :			Tuple									=    ['pixel_values']
          self.assertListEqual(arg_names[:1]   , snake_case_    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			List[Any]									=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*snake_case_    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        def check_hidden_states_output(snake_case_   , snake_case_   , snake_case_    ):
          A_    :			int									=    model_class(snake_case_    )
          model.to(snake_case_    )
          model.eval()

          with torch.no_grad():
            A_    :			List[Any]									=    model(**self._prepare_for_class(snake_case_   , snake_case_    )    )

          A_    :			Dict									=    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
          A_    :			List[str]									=    sum(config.num_block_repeats    ) * 4
          self.assertEqual(len(snake_case_    )   , snake_case_    )

          # EfficientNet'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 // 2, self.model_tester.image_size // 2]   , )

        A_		,			A_    :			str									=    self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          A_    :			Dict									=    True
          check_hidden_states_output(snake_case_   , snake_case_   , snake_case_    )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          A_    :			List[Any]									=    True

          check_hidden_states_output(snake_case_   , snake_case_   , snake_case_    )

      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Optional[Any]									=    self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*snake_case_    )

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

        """simple docstring"""

        pass

      @slow
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          A_    :			List[Any]									=    EfficientNetModel.from_pretrained(snake_case_    )
          self.assertIsNotNone(snake_case_    )

      @is_pipeline_test
      @require_vision
      @slow
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        super().test_pipeline_image_classification()





def 		UpperCAmelCase__						(	):






  """simple docstring"""





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





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



      '''simple docstring'''

      @cached_property
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        return AutoImageProcessor.from_pretrained('google/efficientnet-b7'    ) if is_vision_available() else None

      @slow
      def    lowerCamelCase_     (	self    ):

        """simple docstring"""

        A_    :			Dict									=    EfficientNetForImageClassification.from_pretrained('google/efficientnet-b7'    ).to(snake_case_    )

        A_    :			Tuple									=    self.default_image_processor
        A_    :			Tuple									=    prepare_img()
        A_    :			List[Any]									=    image_processor(images=snake_case_   , return_tensors='pt'    ).to(snake_case_    )

        # forward pass
        with torch.no_grad():
          A_    :			Optional[int]									=    model(**snake_case_    )

        # verify the logits
        A_    :			List[Any]									=    torch.Size((1, 1_0_0_0)    )
        self.assertEqual(outputs.logits.shape   , snake_case_    )

        A_    :			Optional[int]									=    torch.tensor([0.00_01, 0.00_02, 0.00_02]    ).to(snake_case_    )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]   , snake_case_   , atol=1E-4    )    )




import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp

def       A_   (						__a  :					str			,		__a  :					int			,		__a  :					List[Any]			,		__a  :					Dict			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		key.shape[-3:]
				a__					   =		value.shape[-1]
				a__					   =		min(__a			,		__a		)
				a__					   =		query / jnp.sqrt(__a		)

				@functools.partial(jax.checkpoint			,		prevent_cse=__a		)
				def summarize_chunk(__a  :					Optional[Any]			,		__a  :					Tuple			,		__a  :					Dict		):
								a__					   =		jnp.einsum("""...qhd,...khd->...qhk"""			,		__a			,		__a			,		precision=__a		)

								a__					   =		jnp.max(__a			,		axis=-1			,		keepdims=__a		)
								a__					   =		jax.lax.stop_gradient(__a		)
								a__					   =		jnp.exp(attn_weights - max_score		)

								a__					   =		jnp.einsum("""...vhf,...qhv->...qhf"""			,		__a			,		__a			,		precision=__a		)
								a__					   =		jnp.einsum("""...qhk->...qh"""			,		__a		)

								return (exp_values, exp_weights.sum(axis=-1		), max_score)

				def chunk_scanner(__a  :					Tuple		):
								# julienne key array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(key.shape[:-3]		) + [key_chunk_size, num_heads, k_features]			,		)

								# julienne value array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]			,		slice_sizes=list(value.shape[:-3]		) + [key_chunk_size, num_heads, v_features]			,		)

								return summarize_chunk(__a			,		__a			,		__a		)

				a__   ,      a__   ,      a__					   =		jax.lax.map(f=__a			,		xs=jnp.arange(0			,		__a			,		__a		)		)

				a__					   =		jnp.max(__a			,		axis=0			,		keepdims=__a		)
				a__					   =		jnp.exp(chunk_max - global_max		)

				chunk_values *= jnp.expand_dims(__a			,		axis=-1		)
				chunk_weights *= max_diffs

				a__					   =		chunk_values.sum(axis=0		)
				a__					   =		jnp.expand_dims(__a			,		-1		).sum(axis=0		)

				return all_values / all_weights

def       A_   (						__a  :					Tuple			,		__a  :					List[Any]			,		__a  :					int			,		__a  :					List[str]=jax.lax.Precision.HIGHEST			,		__a  :					int = 1_024			,		__a  :					int = 4_096		):

				"""simple docstring"""
				a__   ,      a__   ,      a__					   =		query.shape[-3:]

				def chunk_scanner(__a  :					List[str]			,		__a  :					List[str]		):
								# julienne query array
								a__					   =		jax.lax.dynamic_slice(
								    operand=__a			,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]			,		slice_sizes=list(query.shape[:-3]		) + [min(__a			,		__a		), num_heads, q_features]			,		)

								return (
								    chunk_idx + query_chunk_size,  # unused ignore it
								    _query_chunk_attention(
								        query=__a			,		key=__a			,		value=__a			,		precision=__a			,		key_chunk_size=__a		),
								)

				a__   ,      a__					   =		jax.lax.scan(
				    f=__a			,		init=0			,		xs=__a			,		length=math.ceil(num_q / query_chunk_size		)		)

				return jnp.concatenate(__a			,		axis=-3		)  # fuse the chunked result back

class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       8
							UpperCamelCase__					:   int       =       6_4
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim_head * self.heads
											a__					   =		self.dim_head**-0.5

											# Weights were exported with old names {to_q, to_k, to_v, to_out}
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_q"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_k"""		)
											a__					   =		nn.Dense(a_	,	use_bias=a_	,	dtype=self.dtype	,	name="""to_v"""		)

											a__					   =		nn.Dense(self.query_dim	,	dtype=self.dtype	,	name="""to_out_0"""		)
											a__					   =		nn.Dropout(rate=self.dropout		)


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(a_	,	a_	,	a_	,	dim // head_size		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size * head_size	,	a_	,	dim // head_size		)
											return tensor


							def _a   (				self	,	a_		):
											a__   ,      a__   ,      a__					   =		tensor.shape
											a__					   =		self.heads
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	a_	,	a_		)
											a__					   =		jnp.transpose(a_	,	(0, 2, 1, 3)		)
											a__					   =		tensor.reshape(batch_size // head_size	,	a_	,	dim * head_size		)
											return tensor







							def __call__(				self	,	a_	,	a_=None	,	a_=True		):
											a__					   =		hidden_states if context is None else context

											a__					   =		self.query(a_		)
											a__					   =		self.key(a_		)
											a__					   =		self.value(a_		)

											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)
											a__					   =		self.reshape_heads_to_batch_dim(a_		)

											if self.use_memory_efficient_attention:
															a__					   =		query_states.transpose(1	,	0	,	2		)
															a__					   =		key_states.transpose(1	,	0	,	2		)
															a__					   =		value_states.transpose(1	,	0	,	2		)

															# this if statement create a chunk size for each layer of the unet
															# the chunk size is equal to the query_length dimension of the deepest layer of the unet

															a__					   =		query_states.shape[-3]
															if flatten_latent_dim % 64 == 0:
																			a__					   =		int(flatten_latent_dim / 64		)
															elif flatten_latent_dim % 16 == 0:
																			a__					   =		int(flatten_latent_dim / 16		)
															elif flatten_latent_dim % 4 == 0:
																			a__					   =		int(flatten_latent_dim / 4		)
															else:
																			a__					   =		int(a_		)

															a__					   =		jax_memory_efficient_attention(
															    a_	,	a_	,	a_	,	query_chunk_size=a_	,	key_chunk_size=4_096 * 4		)

															a__					   =		hidden_states.transpose(1	,	0	,	2		)
											else:
															# compute attentions
															a__					   =		jnp.einsum("""b i d, b j d->b i j"""	,	a_	,	a_		)
															a__					   =		attention_scores * self.scale
															a__					   =		nn.softmax(a_	,	axis=2		)

															# attend to values
															a__					   =		jnp.einsum("""b i j, b j d -> b i d"""	,	a_	,	a_		)

											a__					   =		self.reshape_batch_dim_to_heads(a_		)
											a__					   =		self.proj_attn(a_		)
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											# self attention (or cross_attention if only_cross_attention is True)
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											# cross attention
											a__					   =		FlaxAttention(
											    self.dim	,	self.n_heads	,	self.d_head	,	self.dropout	,	self.use_memory_efficient_attention	,	dtype=self.dtype		)
											a__					   =		FlaxFeedForward(dim=self.dim	,	dropout=self.dropout	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.LayerNorm(epsilon=1E-5	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											# self attention
											a__					   =		hidden_states
											if self.only_cross_attention:
															a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											else:
															a__					   =		self.attna(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# cross attention
											a__					   =		hidden_states
											a__					   =		self.attna(self.norma(a_		)	,	a_	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											# feed forward
											a__					   =		hidden_states
											a__					   =		self.ff(self.norma(a_		)	,	deterministic=a_		)
											a__					   =		hidden_states + residual

											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int
							UpperCamelCase__					:   int       =       1
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   bool       =       False
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa
							UpperCamelCase__					:   bool       =       False


							def _a   (				self		):
											a__					   =		nn.GroupNorm(num_groups=32	,	epsilon=1E-5		)

											a__					   =		self.n_heads * self.d_head
											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		[
											    FlaxBasicTransformerBlock(
											        a_	,	self.n_heads	,	self.d_head	,	dropout=self.dropout	,	only_cross_attention=self.only_cross_attention	,	dtype=self.dtype	,	use_memory_efficient_attention=self.use_memory_efficient_attention	,	)
											    for _ in range(self.depth		)
											]

											if self.use_linear_projection:
															a__					   =		nn.Dense(a_	,	dtype=self.dtype		)
											else:
															a__					   =		nn.Conv(
															    a_	,	kernel_size=(1, 1)	,	strides=(1, 1)	,	padding="""VALID"""	,	dtype=self.dtype	,	)

											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_	,	a_=True		):
											a__   ,      a__   ,      a__   ,      a__					   =		hidden_states.shape
											a__					   =		hidden_states
											a__					   =		self.norm(a_		)
											if self.use_linear_projection:
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)
															a__					   =		self.proj_in(a_		)
											else:
															a__					   =		self.proj_in(a_		)
															a__					   =		hidden_states.reshape(a_	,	height * width	,	a_		)

											for transformer_block in self.transformer_blocks:
															a__					   =		transformer_block(a_	,	a_	,	deterministic=a_		)

											if self.use_linear_projection:
															a__					   =		self.proj_out(a_		)
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
											else:
															a__					   =		hidden_states.reshape(a_	,	a_	,	a_	,	a_		)
															a__					   =		self.proj_out(a_		)

											a__					   =		hidden_states + residual
											return self.dropout_layer(a_	,	deterministic=a_		)




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											# The second linear layer needs to be called
											# net_2 for now to match the index of the Sequential layer
											a__					   =		FlaxGEGLU(self.dim	,	self.dropout	,	self.dtype		)
											a__					   =		nn.Dense(self.dim	,	dtype=self.dtype		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.net_a(a_	,	deterministic=a_		)
											a__					   =		self.net_a(a_		)
											return hidden_states




class 			__snake_case   ( nn.Module):







							'''simple docstring'''






							UpperCamelCase__					:   int
							UpperCamelCase__					:   float       =       0.0
							UpperCamelCase__					:   jnp.dtype       =       jnp.floataa


							def _a   (				self		):
											a__					   =		self.dim * 4
											a__					   =		nn.Dense(inner_dim * 2	,	dtype=self.dtype		)
											a__					   =		nn.Dropout(rate=self.dropout		)







							def __call__(				self	,	a_	,	a_=True		):
											a__					   =		self.proj(a_		)
											a__   ,      a__					   =		jnp.split(a_	,	2	,	axis=2		)
											return self.dropout_layer(hidden_linear * nn.gelu(a_		)	,	deterministic=a_		)

