

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


_UpperCAmelCase		:       Dict						    =	logging.get_logger(__name__)


# General docstring
_UpperCAmelCase		:       Union[str, Any]						    =	"""MobileNetV1Config"""

# Base docstring
_UpperCAmelCase		:       Tuple						    =	"""google/mobilenet_v1_1.0_224"""
_UpperCAmelCase		:       Union[str, Any]						    =	[1, 1024, 7, 7]

# Image classification docstring
_UpperCAmelCase		:       Optional[Any]						    =	"""google/mobilenet_v1_1.0_224"""
_UpperCAmelCase		:       Optional[int]						    =	"""tabby, tabby cat"""


_UpperCAmelCase		:       Optional[int]						    =	[
    """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  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__=None				):







							'''simple docstring'''






							snake_case_          =   {}

							if isinstance(__lowerCamelCase       ,							__lowerCamelCase				):
														snake_case_          =   model.mobilenet_va
							else:
														snake_case_          =   model

							snake_case_          =   "MobilenetV1/Conv2d_0/"
							snake_case_          =   backbone.conv_stem.convolution.weight
							snake_case_          =   backbone.conv_stem.normalization.bias
							snake_case_          =   backbone.conv_stem.normalization.weight
							snake_case_          =   backbone.conv_stem.normalization.running_mean
							snake_case_          =   backbone.conv_stem.normalization.running_var

							for i in range(13				):
														snake_case_          =   i + 1
														snake_case_          =   i * 2

														snake_case_          =   backbone.layer[pt_index]
														snake_case_          =   F'''MobilenetV1/Conv2d_{tf_index}_depthwise/'''
														snake_case_          =   pointer.convolution.weight
														snake_case_          =   pointer.normalization.bias
														snake_case_          =   pointer.normalization.weight
														snake_case_          =   pointer.normalization.running_mean
														snake_case_          =   pointer.normalization.running_var

														snake_case_          =   backbone.layer[pt_index + 1]
														snake_case_          =   F'''MobilenetV1/Conv2d_{tf_index}_pointwise/'''
														snake_case_          =   pointer.convolution.weight
														snake_case_          =   pointer.normalization.bias
														snake_case_          =   pointer.normalization.weight
														snake_case_          =   pointer.normalization.running_mean
														snake_case_          =   pointer.normalization.running_var

							if isinstance(__lowerCamelCase       ,							__lowerCamelCase				):
														snake_case_          =   "MobilenetV1/Logits/Conv2d_1c_1x1/"
														snake_case_          =   model.classifier.weight
														snake_case_          =   model.classifier.bias

							return tf_to_pt_map





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__				):







							'''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
							snake_case_          =   tf.train.list_variables(__lowerCamelCase				)
							snake_case_          =   {}
							for name, shape in init_vars:
														logger.info(F'''Loading TF weight {name} with shape {shape}'''				)
														snake_case_          =   tf.train.load_variable(__lowerCamelCase       ,							__lowerCamelCase				)
														snake_case_          =   array

							# Build TF to PyTorch weights loading map
							snake_case_          =   _build_tf_to_pytorch_map(__lowerCamelCase       ,							__lowerCamelCase       ,							__lowerCamelCase				)

							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

														snake_case_          =   tf_weights[name]

														if "depthwise_weights" in name:
																					logger.info('Transposing depthwise'				)
																					snake_case_          =   np.transpose(__lowerCamelCase       ,							(2, 3, 0, 1)				)
														elif "weights" in name:
																					logger.info('Transposing'				)
																					if len(pointer.shape				) == 2:  # copying into linear layer
																												snake_case_          =   array.squeeze().transpose()
																					else:
																												snake_case_          =   np.transpose(__lowerCamelCase       ,							(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}'''				)
														snake_case_          =   torch.from_numpy(__lowerCamelCase				)

														tf_weights.pop(__lowerCamelCase       ,							__lowerCamelCase				)
														tf_weights.pop(name + '/RMSProp'       ,							__lowerCamelCase				)
														tf_weights.pop(name + '/RMSProp_1'       ,							__lowerCamelCase				)
														tf_weights.pop(name + '/ExponentialMovingAverage'       ,							__lowerCamelCase				)

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





def  __lowerCamelCase     (			UpperCamelCase__       ,							UpperCamelCase__				):







							'''simple docstring'''






							snake_case_          =   features.shape[-2:]
							snake_case_          =   conv_layer.stride
							snake_case_          =   conv_layer.kernel_size

							if in_height % stride_height == 0:
														snake_case_          =   max(kernel_height - stride_height       ,							0				)
							else:
														snake_case_          =   max(kernel_height - (in_height % stride_height)       ,							0				)

							if in_width % stride_width == 0:
														snake_case_          =   max(kernel_width - stride_width       ,							0				)
							else:
														snake_case_          =   max(kernel_width - (in_width % stride_width)       ,							0				)

							snake_case_          =   pad_along_width // 2
							snake_case_          =   pad_along_width - pad_left
							snake_case_          =   pad_along_height // 2
							snake_case_          =   pad_along_height - pad_top

							snake_case_          =   (pad_left, pad_right, pad_top, pad_bottom)
							return nn.functional.pad(__lowerCamelCase       ,							__lowerCamelCase       ,							'constant'       ,							0.0				)







class 			lowercase				(	nn.Module      ):




				def __init__( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case = 1				,       snake_case = 1				,       snake_case = False				,       snake_case = True				,       snake_case = True				,       ):
											super().__init__()
											snake_case_          =   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.'''	)

											snake_case_          =   0 if config.tf_padding else int((kernel_size - 1) / 2	)

											snake_case_          =   nn.Convad(
											    in_channels=snake_case__				,       out_channels=snake_case__				,       kernel_size=snake_case__				,       stride=snake_case__				,       padding=snake_case__				,       groups=snake_case__				,       bias=snake_case__				,       padding_mode='zeros'				,       )

											if use_normalization:
																		snake_case_          =   nn.BatchNormad(
																		    num_features=snake_case__				,       eps=config.layer_norm_eps				,       momentum=0.99_97				,       affine=snake_case__				,       track_running_stats=snake_case__				,       )
											else:
																		snake_case_          =   None

											if use_activation:
																		if isinstance(snake_case__				,       snake_case__	):
																									snake_case_          =   ACTaFN[use_activation]
																		elif isinstance(config.hidden_act				,       snake_case__	):
																									snake_case_          =   ACTaFN[config.hidden_act]
																		else:
																									snake_case_          =   config.hidden_act
											else:
																		snake_case_          =   None




				def    a							( self				,       snake_case	):


											if self.config.tf_padding:
																		snake_case_          =   apply_tf_padding(snake_case__				,       self.convolution	)
											snake_case_          =   self.convolution(snake_case__	)
											if self.normalization is not None:
																		snake_case_          =   self.normalization(snake_case__	)
											if self.activation is not None:
																		snake_case_          =   self.activation(snake_case__	)
											return features







class 			lowercase				(	A_      ):
				__SCREAMING_SNAKE_CASE :			List[Any]	      =      MobileNetVaConfig
				__SCREAMING_SNAKE_CASE :			Tuple	      =      load_tf_weights_in_mobilenet_va
				__SCREAMING_SNAKE_CASE :			Any	      =      "mobilenet_v1"
				__SCREAMING_SNAKE_CASE :			Any	      =      "pixel_values"
				__SCREAMING_SNAKE_CASE :			List[str]	      =      False




				def    a							( self				,       snake_case	):


											if isinstance(snake_case__				,       (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(snake_case__				,       nn.BatchNormad	):
																		module.bias.data.zero_()
																		module.weight.data.fill_(1.0	)



_UpperCAmelCase		:       List[str]						    =	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.
"""

_UpperCAmelCase		:       Union[str, Any]						    =	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 			lowercase				(	A_      ):




				def __init__( self				,       snake_case				,       snake_case = True	):
											super().__init__(snake_case__	)
											snake_case_          =   config

											snake_case_          =   32
											snake_case_          =   max(int(depth * config.depth_multiplier	)				,       config.min_depth	)

											snake_case_          =   MobileNetVaConvLayer(
											    snake_case__				,       in_channels=config.num_channels				,       out_channels=snake_case__				,       kernel_size=3				,       stride=2				,       )

											snake_case_          =   [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

											snake_case_          =   nn.ModuleList()
											for i in range(13	):
																		snake_case_          =   out_channels

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

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

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

											snake_case_          =   nn.AdaptiveAvgPoolad((1, 1)	) if add_pooling_layer else None

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




				def    a							( self				,       snake_case	):
											raise NotImplementedError




				@add_start_docstrings_to_model_forward(snake_case__	)
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC				,       output_type=snake_case__				,       config_class=_CONFIG_FOR_DOC				,       modality='vision'				,       expected_output=_EXPECTED_OUTPUT_SHAPE				,       )
				def    a							( self				,       snake_case = None				,       snake_case = None				,       snake_case = None				,       ):
											snake_case_          =   (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											snake_case_          =   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'	)

											snake_case_          =   self.conv_stem(snake_case__	)

											snake_case_          =   () if output_hidden_states else None

											for i, layer_module in enumerate(self.layer	):
																		snake_case_          =   layer_module(snake_case__	)

																		if output_hidden_states:
																									snake_case_          =   all_hidden_states + (hidden_states,)

											snake_case_          =   hidden_states

											if self.pooler is not None:
																		snake_case_          =   torch.flatten(self.pooler(snake_case__	)				,       start_dim=1	)
											else:
																		snake_case_          =   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=snake_case__				,       pooler_output=snake_case__				,       hidden_states=snake_case__				,       )







@add_start_docstrings(
    '''\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '''     ,  A_     ,  )
class 			lowercase				(	A_      ):




				def __init__( self				,       snake_case	):
											super().__init__(snake_case__	)

											snake_case_          =   config.num_labels
											snake_case_          =   MobileNetVaModel(snake_case__	)

											snake_case_          =   self.mobilenet_va.layer[-1].convolution.out_channels

											# Classifier head
											snake_case_          =   nn.Dropout(config.classifier_dropout_prob				,       inplace=snake_case__	)
											snake_case_          =   nn.Linear(snake_case__				,       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(snake_case__	)
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT				,       output_type=snake_case__				,       config_class=_CONFIG_FOR_DOC				,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,       )
				def    a							( self				,       snake_case = None				,       snake_case = None				,       snake_case = None				,       snake_case = None				,       ):
											snake_case_          =   return_dict if return_dict is not None else self.config.use_return_dict

											snake_case_          =   self.mobilenet_va(snake_case__				,       output_hidden_states=snake_case__				,       return_dict=snake_case__	)

											snake_case_          =   outputs.pooler_output if return_dict else outputs[1]

											snake_case_          =   self.classifier(self.dropout(snake_case__	)	)

											snake_case_          =   None
											if labels is not None:
																		if self.config.problem_type is None:
																									if self.num_labels == 1:
																																snake_case_          =   "regression"
																									elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																																snake_case_          =   "single_label_classification"
																									else:
																																snake_case_          =   "multi_label_classification"

																		if self.config.problem_type == "regression":
																									snake_case_          =   MSELoss()
																									if self.num_labels == 1:
																																snake_case_          =   loss_fct(logits.squeeze()				,       labels.squeeze()	)
																									else:
																																snake_case_          =   loss_fct(snake_case__				,       snake_case__	)
																		elif self.config.problem_type == "single_label_classification":
																									snake_case_          =   CrossEntropyLoss()
																									snake_case_          =   loss_fct(logits.view(-1				,       self.num_labels	)				,       labels.view(-1	)	)
																		elif self.config.problem_type == "multi_label_classification":
																									snake_case_          =   BCEWithLogitsLoss()
																									snake_case_          =   loss_fct(snake_case__				,       snake_case__	)

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

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







def  __lowerCamelCase     (			UpperCamelCase__				):







							'''simple docstring'''






							snake_case_          =   1
							snake_case_          =   2
							while i * i <= n:
														snake_case_          =   0
														while n % i == 0:
																					n //= i
																					multiplicity += 1
														n_divisors *= multiplicity + 1
														i += 1
							if n > 1:
														n_divisors *= 2
							return n_divisors





def  __lowerCamelCase     (			):







							'''simple docstring'''






							snake_case_          =   1
							snake_case_          =   1

							while True:
														i += 1
														t_num += i

														if count_divisors(UpperCamelCase__				) > 500:
																					break

							return t_num


if __name__ == "__main__":
					print(solution())




