


"""simple docstring"""







import unittest

from transformers import PegasusConfig, PegasusTokenizer, is_flax_available
from transformers.testing_utils import require_flax, slow

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


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 __lowercase    =				"""platform"""
 import jax
 import jax.numpy as jnp
 import numpy as np

 from transformers import FlaxPegasusForConditionalGeneration, FlaxPegasusModel






@require_flax
class 						_A  :







    """simple docstring"""




    UpperCAmelCase	:							List[Any]     =	PegasusConfig
    UpperCAmelCase	:							Dict     =	{}
    UpperCAmelCase	:							Tuple     =	"""gelu"""






    def __init__( self    :   Any    ,			__UpperCAmelCase    :   Dict    ,			__UpperCAmelCase    :   int=13    ,			__UpperCAmelCase    :   List[str]=7    ,			__UpperCAmelCase    :   Dict=True    ,			__UpperCAmelCase    :   List[Any]=False    ,			__UpperCAmelCase    :   int=99    ,			__UpperCAmelCase    :   Tuple=32    ,			__UpperCAmelCase    :   List[Any]=5    ,			__UpperCAmelCase    :   Optional[int]=4    ,			__UpperCAmelCase    :   Tuple=37    ,			__UpperCAmelCase    :   List[Any]=0.1    ,			__UpperCAmelCase    :   Dict=0.1    ,			__UpperCAmelCase    :   Dict=20    ,			__UpperCAmelCase    :   Optional[Any]=2    ,			__UpperCAmelCase    :   Optional[Any]=1    ,			__UpperCAmelCase    :   List[Any]=0    ,			):
      a						:		Any											=      parent
      a						:		Union[str, Any]											=      batch_size
      a						:		Dict											=      seq_length
      a						:		Dict											=      is_training
      a						:		Optional[Any]											=      use_labels
      a						:		Tuple											=      vocab_size
      a						:		int											=      hidden_size
      a						:		List[Any]											=      num_hidden_layers
      a						:		int											=      num_attention_heads
      a						:		List[Any]											=      intermediate_size

      a						:		Any											=      hidden_dropout_prob
      a						:		str											=      attention_probs_dropout_prob
      a						:		List[Any]											=      max_position_embeddings
      a						:		Optional[Any]											=      eos_token_id
      a						:		Optional[int]											=      pad_token_id
      a						:		Dict											=      bos_token_id






    def      __snake_case							( self    :   Dict):
      a						:		Tuple											=      ids_tensor([self.batch_size, self.seq_length - 1]    ,			self.vocab_size).clip(3    ,			self.vocab_size)
      a						:		Optional[Any]											=      np.expand_dims(np.array([self.eos_token_id] * self.batch_size)    ,			1)
      a						:		Any											=      np.concatenate([input_ids, eos_tensor]    ,			axis=1)

      a						:		Optional[Any]											=      ids_tensor([self.batch_size, self.seq_length]    ,			self.vocab_size)

      a						:		Optional[int]											=      self.config_cls(
          vocab_size=self.vocab_size    ,			d_model=self.hidden_size    ,			encoder_layers=self.num_hidden_layers    ,			decoder_layers=self.num_hidden_layers    ,			encoder_attention_heads=self.num_attention_heads    ,			decoder_attention_heads=self.num_attention_heads    ,			encoder_ffn_dim=self.intermediate_size    ,			decoder_ffn_dim=self.intermediate_size    ,			dropout=self.hidden_dropout_prob    ,			attention_dropout=self.attention_probs_dropout_prob    ,			max_position_embeddings=self.max_position_embeddings    ,			eos_token_ids=[2]    ,			bos_token_id=self.bos_token_id    ,			pad_token_id=self.pad_token_id    ,			decoder_start_token_id=self.pad_token_id    ,			**self.config_updates    ,			)
      a						:		Optional[int]											=      prepare_pegasus_inputs_dict(__UpperCAmelCase    ,			__UpperCAmelCase    ,			__UpperCAmelCase)
      return config, inputs_dict






    def      __snake_case							( self    :   Union[str, Any]    ,			__UpperCAmelCase    :   Dict    ,			__UpperCAmelCase    :   List[str]    ,			__UpperCAmelCase    :   Any):
      a						:		int											=      20
      a						:		Tuple											=      model_class_name(__UpperCAmelCase)

      a						:		Optional[int]											=      model.encode(inputs_dict["input_ids"])

      a						:		Union[str, Any]											=      (
          inputs_dict["decoder_input_ids"],
          inputs_dict["decoder_attention_mask"],
      )

      a						:		Optional[int]											=      model.init_cache(decoder_input_ids.shape[0]    ,			__UpperCAmelCase    ,			__UpperCAmelCase)
      a						:		Any											=      jnp.ones((decoder_input_ids.shape[0], max_decoder_length)    ,			dtype="i4")

      a						:		List[str]											=      jnp.broadcast_to(
          jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]    ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)    ,			)
      a						:		Tuple											=      model.decode(
          decoder_input_ids[:, :-1]    ,			__UpperCAmelCase    ,			decoder_attention_mask=__UpperCAmelCase    ,			past_key_values=__UpperCAmelCase    ,			decoder_position_ids=__UpperCAmelCase    ,			)

      a						:		int											=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    ,			dtype="i4")
      a						:		Any											=      model.decode(
          decoder_input_ids[:, -1:]    ,			__UpperCAmelCase    ,			decoder_attention_mask=__UpperCAmelCase    ,			past_key_values=outputs_cache.past_key_values    ,			decoder_position_ids=__UpperCAmelCase    ,			)

      a						:		Any											=      model.decode(__UpperCAmelCase    ,			__UpperCAmelCase)

      a						:		List[Any]											=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
      self.parent.assertTrue(diff < 1e-3    ,			msg=f'''Max diff is {diff}''')






    def      __snake_case							( self    :   Union[str, Any]    ,			__UpperCAmelCase    :   str    ,			__UpperCAmelCase    :   Union[str, Any]    ,			__UpperCAmelCase    :   Dict):
      a						:		int											=      20
      a						:		List[Any]											=      model_class_name(__UpperCAmelCase)

      a						:		Tuple											=      model.encode(inputs_dict["input_ids"])

      a						:		Tuple											=      (
          inputs_dict["decoder_input_ids"],
          inputs_dict["decoder_attention_mask"],
      )

      a						:		int											=      jnp.concatenate(
          [
              decoder_attention_mask,
              jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
          ]    ,			axis=-1    ,			)

      a						:		Any											=      model.init_cache(decoder_input_ids.shape[0]    ,			__UpperCAmelCase    ,			__UpperCAmelCase)
      a						:		Optional[Any]											=      jnp.broadcast_to(
          jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]    ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)    ,			)

      a						:		int											=      model.decode(
          decoder_input_ids[:, :-1]    ,			__UpperCAmelCase    ,			decoder_attention_mask=__UpperCAmelCase    ,			past_key_values=__UpperCAmelCase    ,			decoder_position_ids=__UpperCAmelCase    ,			)
      a						:		Union[str, Any]											=      jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]    ,			dtype="i4")
      a						:		Tuple											=      model.decode(
          decoder_input_ids[:, -1:]    ,			__UpperCAmelCase    ,			past_key_values=outputs_cache.past_key_values    ,			decoder_attention_mask=__UpperCAmelCase    ,			decoder_position_ids=__UpperCAmelCase    ,			)

      a						:		Tuple											=      model.decode(__UpperCAmelCase    ,			__UpperCAmelCase    ,			decoder_attention_mask=__UpperCAmelCase)

      a						:		List[str]											=      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
      self.parent.assertTrue(diff < 1e-3    ,			msg=f'''Max diff is {diff}''')






def 					lowercase      (    A_  ,       A_  ,       A_  ,       A_=None  ,       A_=None  ,       )->				str:
  '''simple docstring'''







  if attention_mask is None:
    a						:		str											=      np.not_equal(A_  ,       config.pad_token_id	).astype(np.inta	)
  if decoder_attention_mask is None:
    a						:		List[str]											=      np.concatenate(
        [
            np.ones(decoder_input_ids[:, :1].shape  ,       dtype=np.inta	),
            np.not_equal(decoder_input_ids[:, 1:]  ,       config.pad_token_id	).astype(np.inta	),
        ]  ,       axis=-1  ,       )
  return {
      "input_ids": input_ids,
      "decoder_input_ids": decoder_input_ids,
      "attention_mask": attention_mask,
      "decoder_attention_mask": decoder_attention_mask,
  }







@require_flax
class 						_A  (  _a   ,unittest.TestCase  ):







    """simple docstring"""




    UpperCAmelCase	:							Tuple     =	(
        (
            FlaxPegasusForConditionalGeneration,
            FlaxPegasusModel,
        )
        if is_flax_available()
        else ()
    )
    UpperCAmelCase	:							Union[str, Any]     =	(FlaxPegasusForConditionalGeneration,) if is_flax_available() else ()
    UpperCAmelCase	:							Tuple     =	True
    UpperCAmelCase	:							Optional[Any]     =	False
    UpperCAmelCase	:							Union[str, Any]     =	False
    UpperCAmelCase	:							Union[str, Any]     =	False






    def      __snake_case							( self    :   Union[str, Any]):
      a						:		Optional[Any]											=      FlaxPegasusModelTester(self)
      a						:		List[str]											=      ConfigTester(self    ,			config_class=__UpperCAmelCase)






    def      __snake_case							( self    :   Any):
      self.config_tester.run_common_tests()






    def      __snake_case							( self    :   Union[str, Any]):
      a						:		Optional[int]											=      self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward(__UpperCAmelCase    ,			__UpperCAmelCase    ,			__UpperCAmelCase)






    def      __snake_case							( self    :   List[str]):
      a						:		Any											=      self.model_tester.prepare_config_and_inputs_for_common()
      for model_class in self.all_model_classes:
        self.model_tester.check_use_cache_forward_with_attn_mask(__UpperCAmelCase    ,			__UpperCAmelCase    ,			__UpperCAmelCase)






    def      __snake_case							( self    :   Optional[Any]):
      a						:		Optional[int]											=      self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__):
          a						:		List[str]											=      self._prepare_for_class(__UpperCAmelCase    ,			__UpperCAmelCase)
          a						:		List[Any]											=      model_class(__UpperCAmelCase)

          @jax.jit
          def encode_jitted(__UpperCAmelCase    :   Tuple    ,			__UpperCAmelCase    :   Optional[Any]=None    ,			**__UpperCAmelCase    :   Any):
            return model.encode(input_ids=__UpperCAmelCase    ,			attention_mask=__UpperCAmelCase)

          with self.subTest("JIT Enabled"):
            a						:		Optional[Any]											=      encode_jitted(**__UpperCAmelCase).to_tuple()

          with self.subTest("JIT Disabled"):
            with jax.disable_jit():
              a						:		Optional[Any]											=      encode_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      __snake_case							( self    :   Tuple):
      a						:		Tuple											=      self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__):
          a						:		Union[str, Any]											=      model_class(__UpperCAmelCase)
          a						:		List[str]											=      model.encode(inputs_dict["input_ids"]    ,			inputs_dict["attention_mask"])

          a						:		Optional[int]											=      {
              "decoder_input_ids": inputs_dict["decoder_input_ids"],
              "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
              "encoder_outputs": encoder_outputs,
          }

          @jax.jit
          def decode_jitted(__UpperCAmelCase    :   Optional[int]    ,			__UpperCAmelCase    :   Optional[Any]    ,			__UpperCAmelCase    :   List[str]):
            return model.decode(
                decoder_input_ids=__UpperCAmelCase    ,			decoder_attention_mask=__UpperCAmelCase    ,			encoder_outputs=__UpperCAmelCase    ,			)

          with self.subTest("JIT Enabled"):
            a						:		str											=      decode_jitted(**__UpperCAmelCase).to_tuple()

          with self.subTest("JIT Disabled"):
            with jax.disable_jit():
              a						:		List[str]											=      decode_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)






    @slow
    def      __snake_case							( self    :   Dict):
      for model_class_name in self.all_model_classes:
        a						:		Dict											=      model_class_name.from_pretrained("google/pegasus-large"    ,			from_pt=__UpperCAmelCase)
        a						:		str											=      np.ones((1, 1))
        a						:		int											=      model(__UpperCAmelCase)
        self.assertIsNotNone(__UpperCAmelCase)






    @slow
    def      __snake_case							( self    :   Any):
      a						:		str											=      FlaxPegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
      a						:		List[Any]											=      PegasusTokenizer.from_pretrained("google/pegasus-xsum")

      a						:		str											=      [
          " PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow.",
          " The London trio are up for best UK act and best album, as well as getting two nominations in the best song category.\"We got told like this morning 'Oh I think you're nominated'\", said Dappy.\"And I was like 'Oh yeah, which one?' And now we've got nominated for four awards. I mean, wow!\"Bandmate Fazer added: \"We thought it's best of us to come down and mingle with everyone and say hello to the cameras. And now we find we've got four nominations.\"The band have two shots at the best song prize, getting the nod for their Tynchy Stryder collaboration Number One, and single Strong Again.Their album Uncle B will also go up against records by the likes of Beyonce and Kanye West.N-Dubz picked up the best newcomer Mobo in 2007, but female member Tulisa said they wouldn't be too disappointed if they didn't win this time around.\"At the end of the day we're grateful to be where we are in our careers.\"If it don't happen then it don't happen - live to fight another day and keep on making albums and hits for the fans.\"Dappy also revealed they could be performing live several times on the night.The group will be doing Number One and also a possible rendition of the War Child single, I Got Soul.The charity song is a  re-working of The Killers' All These Things That I've Done and is set to feature artists like Chipmunk, Ironik and Pixie Lott.This year's Mobos will be held outside of London for the first time, in Glasgow on 30 September.N-Dubz said they were looking forward to performing for their Scottish fans and boasted about their recent shows north of the border.\"We just done Edinburgh the other day,\" said Dappy.\"We smashed up an N-Dubz show over there. We done Aberdeen about three or four months ago - we smashed up that show over there! Everywhere we go we smash it up!\" ",
      ]

      a						:		str											=      [
          "California's largest electricity provider has turned off power to hundreds of thousands of customers.",
          "Pop group N-Dubz have revealed they were surprised to get four nominations for this year's Mobo Awards.",
      ]

      a						:		Any											=      tokenizer(__UpperCAmelCase    ,			return_tensors="np"    ,			truncation=__UpperCAmelCase    ,			max_length=512    ,			padding=__UpperCAmelCase)
      a						:		Union[str, Any]											=      model.generate(**__UpperCAmelCase    ,			num_beams=2).sequences
      a						:		Tuple											=      tokenizer.batch_decode(__UpperCAmelCase    ,			skip_special_tokens=__UpperCAmelCase)
      assert tgt_text == decoded










"""simple docstring"""







from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


__lowercase    =				logging.get_logger(__name__)


# General docstring
__lowercase    =				"""MobileNetV1Config"""

# Base docstring
__lowercase    =				"""google/mobilenet_v1_1.0_224"""
__lowercase    =				[1, 1024, 7, 7]

# Image classification docstring
__lowercase    =				"""google/mobilenet_v1_1.0_224"""
__lowercase    =				"""tabby, tabby cat"""


__lowercase    =				[
    """google/mobilenet_v1_1.0_224""",
    """google/mobilenet_v1_0.75_192""",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]


def 					lowercase      (    A_  ,       A_  ,       A_=None	)->				int:
  '''simple docstring'''







  a						:		Union[str, Any]											=      {}

  if isinstance(A_  ,       A_	):
    a						:		Tuple											=      model.mobilenet_va
  else:
    a						:		Optional[int]											=      model

  a						:		Any											=      "MobilenetV1/Conv2d_0/"
  a						:		List[Any]											=      backbone.conv_stem.convolution.weight
  a						:		Tuple											=      backbone.conv_stem.normalization.bias
  a						:		int											=      backbone.conv_stem.normalization.weight
  a						:		Optional[int]											=      backbone.conv_stem.normalization.running_mean
  a						:		Dict											=      backbone.conv_stem.normalization.running_var

  for i in range(13	):
    a						:		Dict											=      i + 1
    a						:		str											=      i * 2

    a						:		str											=      backbone.layer[pt_index]
    a						:		Tuple											=      F'''MobilenetV1/Conv2d_{tf_index}_depthwise/'''
    a						:		Union[str, Any]											=      pointer.convolution.weight
    a						:		Optional[int]											=      pointer.normalization.bias
    a						:		Any											=      pointer.normalization.weight
    a						:		Optional[int]											=      pointer.normalization.running_mean
    a						:		Dict											=      pointer.normalization.running_var

    a						:		Dict											=      backbone.layer[pt_index + 1]
    a						:		List[str]											=      F'''MobilenetV1/Conv2d_{tf_index}_pointwise/'''
    a						:		Dict											=      pointer.convolution.weight
    a						:		Any											=      pointer.normalization.bias
    a						:		List[Any]											=      pointer.normalization.weight
    a						:		Tuple											=      pointer.normalization.running_mean
    a						:		List[str]											=      pointer.normalization.running_var

  if isinstance(A_  ,       A_	):
    a						:		Dict											=      "MobilenetV1/Logits/Conv2d_1c_1x1/"
    a						:		List[Any]											=      model.classifier.weight
    a						:		Optional[int]											=      model.classifier.bias

  return tf_to_pt_map


def 					lowercase      (    A_  ,       A_  ,       A_	)->				Optional[int]:
  '''simple docstring'''







  try:
    import numpy as np
    import tensorflow as tf
  except ImportError:
    logger.error(
        "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
        "https://www.tensorflow.org/install/ for installation instructions."	)
    raise

  # Load weights from TF model
  a						:		List[str]											=      tf.train.list_variables(A_	)
  a						:		Dict											=      {}
  for name, shape in init_vars:
    logger.info(F'''Loading TF weight {name} with shape {shape}'''	)
    a						:		List[str]											=      tf.train.load_variable(A_  ,       A_	)
    a						:		int											=      array

  # Build TF to PyTorch weights loading map
  a						:		int											=      _build_tf_to_pytorch_map(A_  ,       A_  ,       A_	)

  for name, pointer in tf_to_pt_map.items():
    logger.info(F'''Importing {name}'''	)
    if name not in tf_weights:
      logger.info(F'''{name} not in tf pre-trained weights, skipping'''	)
      continue

    a						:		Tuple											=      tf_weights[name]

    if "depthwise_weights" in name:
      logger.info("Transposing depthwise"	)
      a						:		Optional[Any]											=      np.transpose(A_  ,       (2, 3, 0, 1)	)
    elif "weights" in name:
      logger.info("Transposing"	)
      if len(pointer.shape	) == 2:  # copying into linear layer
        a						:		List[str]											=      array.squeeze().transpose()
      else:
        a						:		List[str]											=      np.transpose(A_  ,       (3, 2, 0, 1)	)

    if pointer.shape != array.shape:
      raise ValueError(F'''Pointer shape {pointer.shape} and array shape {array.shape} mismatched'''	)

    logger.info(F'''Initialize PyTorch weight {name} {array.shape}'''	)
    a						:		Any											=      torch.from_numpy(A_	)

    tf_weights.pop(A_  ,       A_	)
    tf_weights.pop(name + "/RMSProp"  ,       A_	)
    tf_weights.pop(name + "/RMSProp_1"  ,       A_	)
    tf_weights.pop(name + "/ExponentialMovingAverage"  ,       A_	)

  logger.info(F'''Weights not copied to PyTorch model: {", ".join(tf_weights.keys()	)}'''	)
  return model


def 					lowercase      (    A_  ,       A_	)->				torch.Tensor:
  '''simple docstring'''







  a							,		a						:		str											=      features.shape[-2:]
  a							,		a						:		Tuple											=      conv_layer.stride
  a							,		a						:		Union[str, Any]											=      conv_layer.kernel_size

  if in_height % stride_height == 0:
    a						:		Union[str, Any]											=      max(kernel_height - stride_height  ,       0	)
  else:
    a						:		Tuple											=      max(kernel_height - (in_height % stride_height)  ,       0	)

  if in_width % stride_width == 0:
    a						:		Optional[int]											=      max(kernel_width - stride_width  ,       0	)
  else:
    a						:		int											=      max(kernel_width - (in_width % stride_width)  ,       0	)

  a						:		List[Any]											=      pad_along_width // 2
  a						:		List[str]											=      pad_along_width - pad_left
  a						:		str											=      pad_along_height // 2
  a						:		Any											=      pad_along_height - pad_top

  a						:		Tuple											=      (pad_left, pad_right, pad_top, pad_bottom)
  return nn.functional.pad(A_  ,       A_  ,       "constant"  ,       0.0	)






class 						_A  (  nn.Module  ):







    """simple docstring"""




    def __init__( self    :   Dict    ,			__UpperCAmelCase    :   MobileNetVaConfig    ,			__UpperCAmelCase    :   int    ,			__UpperCAmelCase    :   int    ,			__UpperCAmelCase    :   int    ,			__UpperCAmelCase    :   Optional[int] = 1    ,			__UpperCAmelCase    :   Optional[int] = 1    ,			__UpperCAmelCase    :   bool = False    ,			__UpperCAmelCase    :   Optional[bool] = True    ,			__UpperCAmelCase    :   Optional[bool or str] = True    ,			):
      super().__init__()
      a						:		Optional[Any]											=      config

      if in_channels % groups != 0:
        raise ValueError(f'''Input channels ({in_channels}) are not divisible by {groups} groups.''')
      if out_channels % groups != 0:
        raise ValueError(f'''Output channels ({out_channels}) are not divisible by {groups} groups.''')

      a						:		Dict											=      0 if config.tf_padding else int((kernel_size - 1) / 2)

      a						:		List[str]											=      nn.Convad(
          in_channels=__UpperCAmelCase    ,			out_channels=__UpperCAmelCase    ,			kernel_size=__UpperCAmelCase    ,			stride=__UpperCAmelCase    ,			padding=__UpperCAmelCase    ,			groups=__UpperCAmelCase    ,			bias=__UpperCAmelCase    ,			padding_mode="zeros"    ,			)

      if use_normalization:
        a						:		str											=      nn.BatchNormad(
            num_features=__UpperCAmelCase    ,			eps=config.layer_norm_eps    ,			momentum=0.9_997    ,			affine=__UpperCAmelCase    ,			track_running_stats=__UpperCAmelCase    ,			)
      else:
        a						:		List[str]											=      None

      if use_activation:
        if isinstance(__UpperCAmelCase    ,			__UpperCAmelCase):
          a						:		Optional[int]											=      ACTaFN[use_activation]
        elif isinstance(config.hidden_act    ,			__UpperCAmelCase):
          a						:		Optional[Any]											=      ACTaFN[config.hidden_act]
        else:
          a						:		str											=      config.hidden_act
      else:
        a						:		Dict											=      None






    def      __snake_case							( self    :   Tuple    ,			__UpperCAmelCase    :   torch.Tensor):
      if self.config.tf_padding:
        a						:		Tuple											=      apply_tf_padding(__UpperCAmelCase    ,			self.convolution)
      a						:		Dict											=      self.convolution(__UpperCAmelCase)
      if self.normalization is not None:
        a						:		List[str]											=      self.normalization(__UpperCAmelCase)
      if self.activation is not None:
        a						:		Optional[Any]											=      self.activation(__UpperCAmelCase)
      return features






class 						_A  (  _a  ):







    """simple docstring"""




    UpperCAmelCase	:							Union[str, Any]     =	MobileNetVaConfig
    UpperCAmelCase	:							str     =	load_tf_weights_in_mobilenet_va
    UpperCAmelCase	:							List[str]     =	"""mobilenet_v1"""
    UpperCAmelCase	:							Dict     =	"""pixel_values"""
    UpperCAmelCase	:							str     =	False






    def      __snake_case							( self    :   str    ,			__UpperCAmelCase    :   Union[nn.Linear, nn.Convad]):
      if isinstance(__UpperCAmelCase    ,			(nn.Linear, nn.Convad)):
        module.weight.data.normal_(mean=0.0    ,			std=self.config.initializer_range)
        if module.bias is not None:
          module.bias.data.zero_()
      elif isinstance(__UpperCAmelCase    ,			nn.BatchNormad):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)






__lowercase    =				R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

__lowercase    =				R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`MobileNetV1ImageProcessor.__call__`] for details.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""






@add_start_docstrings(
    """The bare MobileNetV1 model outputting raw hidden-states without any specific head on top."""   ,_a   ,)
class 						_A  (  _a  ):







    """simple docstring"""




    def __init__( self    :   Optional[int]    ,			__UpperCAmelCase    :   MobileNetVaConfig    ,			__UpperCAmelCase    :   bool = True):
      super().__init__(__UpperCAmelCase)
      a						:		List[str]											=      config

      a						:		List[Any]											=      32
      a						:		Union[str, Any]											=      max(int(depth * config.depth_multiplier)    ,			config.min_depth)

      a						:		List[str]											=      MobileNetVaConvLayer(
          __UpperCAmelCase    ,			in_channels=config.num_channels    ,			out_channels=__UpperCAmelCase    ,			kernel_size=3    ,			stride=2    ,			)

      a						:		int											=      [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

      a						:		List[Any]											=      nn.ModuleList()
      for i in range(13):
        a						:		Optional[Any]											=      out_channels

        if strides[i] == 2 or i == 0:
          depth *= 2
          a						:		List[str]											=      max(int(depth * config.depth_multiplier)    ,			config.min_depth)

        self.layer.append(
            MobileNetVaConvLayer(
                __UpperCAmelCase    ,			in_channels=__UpperCAmelCase    ,			out_channels=__UpperCAmelCase    ,			kernel_size=3    ,			stride=strides[i]    ,			groups=__UpperCAmelCase    ,			))

        self.layer.append(
            MobileNetVaConvLayer(
                __UpperCAmelCase    ,			in_channels=__UpperCAmelCase    ,			out_channels=__UpperCAmelCase    ,			kernel_size=1    ,			))

      a						:		Optional[Any]											=      nn.AdaptiveAvgPoolad((1, 1)) if add_pooling_layer else None

      # Initialize weights and apply final processing
      self.post_init()






    def      __snake_case							( self    :   Optional[Any]    ,			__UpperCAmelCase    :   Any):
      raise NotImplementedError






    @add_start_docstrings_to_model_forward(__UpperCAmelCase)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC    ,			output_type=__UpperCAmelCase    ,			config_class=_CONFIG_FOR_DOC    ,			modality="vision"    ,			expected_output=_EXPECTED_OUTPUT_SHAPE    ,			)
    def      __snake_case							( self    :   List[Any]    ,			__UpperCAmelCase    :   Optional[torch.Tensor] = None    ,			__UpperCAmelCase    :   Optional[bool] = None    ,			__UpperCAmelCase    :   Optional[bool] = None    ,			):
      a						:		str											=      (
          output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
      )
      a						:		Optional[Any]											=      return_dict if return_dict is not None else self.config.use_return_dict

      if pixel_values is None:
        raise ValueError("You have to specify pixel_values")

      a						:		Optional[int]											=      self.conv_stem(__UpperCAmelCase)

      a						:		Tuple											=      () if output_hidden_states else None

      for i, layer_module in enumerate(self.layer):
        a						:		Union[str, Any]											=      layer_module(__UpperCAmelCase)

        if output_hidden_states:
          a						:		str											=      all_hidden_states + (hidden_states,)

      a						:		Dict											=      hidden_states

      if self.pooler is not None:
        a						:		Optional[int]											=      torch.flatten(self.pooler(__UpperCAmelCase)    ,			start_dim=1)
      else:
        a						:		List[str]											=      None

      if not return_dict:
        return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None)

      return BaseModelOutputWithPoolingAndNoAttention(
          last_hidden_state=__UpperCAmelCase    ,			pooler_output=__UpperCAmelCase    ,			hidden_states=__UpperCAmelCase    ,			)






@add_start_docstrings(
    """
    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """   ,_a   ,)
class 						_A  (  _a  ):







    """simple docstring"""




    def __init__( self    :   int    ,			__UpperCAmelCase    :   MobileNetVaConfig):
      super().__init__(__UpperCAmelCase)

      a						:		Dict											=      config.num_labels
      a						:		Union[str, Any]											=      MobileNetVaModel(__UpperCAmelCase)

      a						:		List[Any]											=      self.mobilenet_va.layer[-1].convolution.out_channels

      # Classifier head
      a						:		Union[str, Any]											=      nn.Dropout(config.classifier_dropout_prob    ,			inplace=__UpperCAmelCase)
      a						:		Optional[int]											=      nn.Linear(__UpperCAmelCase    ,			config.num_labels) if config.num_labels > 0 else nn.Identity()

      # Initialize weights and apply final processing
      self.post_init()






    @add_start_docstrings_to_model_forward(__UpperCAmelCase)
    @add_code_sample_docstrings(
        checkpoint=_IMAGE_CLASS_CHECKPOINT    ,			output_type=__UpperCAmelCase    ,			config_class=_CONFIG_FOR_DOC    ,			expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,			)
    def      __snake_case							( self    :   Dict    ,			__UpperCAmelCase    :   Optional[torch.Tensor] = None    ,			__UpperCAmelCase    :   Optional[bool] = None    ,			__UpperCAmelCase    :   Optional[torch.Tensor] = None    ,			__UpperCAmelCase    :   Optional[bool] = None    ,			):
      a						:		Union[str, Any]											=      return_dict if return_dict is not None else self.config.use_return_dict

      a						:		Dict											=      self.mobilenet_va(__UpperCAmelCase    ,			output_hidden_states=__UpperCAmelCase    ,			return_dict=__UpperCAmelCase)

      a						:		List[Any]											=      outputs.pooler_output if return_dict else outputs[1]

      a						:		List[Any]											=      self.classifier(self.dropout(__UpperCAmelCase))

      a						:		Any											=      None
      if labels is not None:
        if self.config.problem_type is None:
          if self.num_labels == 1:
            a						:		List[str]											=      "regression"
          elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
            a						:		Any											=      "single_label_classification"
          else:
            a						:		Optional[Any]											=      "multi_label_classification"

        if self.config.problem_type == "regression":
          a						:		Union[str, Any]											=      MSELoss()
          if self.num_labels == 1:
            a						:		Tuple											=      loss_fct(logits.squeeze()    ,			labels.squeeze())
          else:
            a						:		Dict											=      loss_fct(__UpperCAmelCase    ,			__UpperCAmelCase)
        elif self.config.problem_type == "single_label_classification":
          a						:		Optional[Any]											=      CrossEntropyLoss()
          a						:		Tuple											=      loss_fct(logits.view(-1    ,			self.num_labels)    ,			labels.view(-1))
        elif self.config.problem_type == "multi_label_classification":
          a						:		Dict											=      BCEWithLogitsLoss()
          a						:		List[Any]											=      loss_fct(__UpperCAmelCase    ,			__UpperCAmelCase)

      if not return_dict:
        a						:		int											=      (logits,) + outputs[2:]
        return ((loss,) + output) if loss is not None else output

      return ImageClassifierOutputWithNoAttention(
          loss=__UpperCAmelCase    ,			logits=__UpperCAmelCase    ,			hidden_states=outputs.hidden_states    ,			)






