





import collections.abc
from typing import Optional, Tuple, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, 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_poolformer import PoolFormerConfig


_lowerCamelCase :    str  							= logging.get_logger(__name__)

# General docstring
_lowerCamelCase :    Tuple  							= '''PoolFormerConfig'''

# Base docstring
_lowerCamelCase :    List[Any]  							= '''sail/poolformer_s12'''
_lowerCamelCase :    int  							= [1, 5_1_2, 7, 7]

# Image classification docstring
_lowerCamelCase :    int  							= '''sail/poolformer_s12'''
_lowerCamelCase :    Union[str, Any]  							= '''tabby, tabby cat'''

_lowerCamelCase :    Tuple  							= [
    '''sail/poolformer_s12''',
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]





def _a (  SCREAMING_SNAKE_CASE__      :  Union[str, Any]			,				SCREAMING_SNAKE_CASE__      :  float = 0.0			,				SCREAMING_SNAKE_CASE__      :  bool = False       )	->    int:

							'''simple docstring'''






							if drop_prob == 0.0 or not training:
														return input
							SCREAMING_SNAKE_CASE__       :      Any     =						1 - drop_prob
							SCREAMING_SNAKE_CASE__       :      Any     =						(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						keep_prob + torch.rand(SCREAMING_SNAKE_CASE__			,				dtype=input.dtype			,				device=input.device       )
							random_tensor.floor_()  # binarize
							SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						input.div(SCREAMING_SNAKE_CASE__       ) * random_tensor
							return output


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					List[Any],			_UpperCAmelCase     :					Optional[float] = None						)				->			None:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      int     =						drop_prob






							def  A_     ( self     :					str,			_UpperCAmelCase     :					torch.Tensor						)				->			torch.Tensor:





														"""simple docstring"""
														return drop_path(_UpperCAmelCase,			self.drop_prob,			self.training						)






							def  A_     ( self     :					Tuple						)				->			str:





														"""simple docstring"""
														return "p={}".format(self.drop_prob						)


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					Optional[Any],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					str=None						)				->			List[Any]:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						patch_size if isinstance(_UpperCAmelCase,			collections.abc.Iterable						) else (patch_size, patch_size)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						stride if isinstance(_UpperCAmelCase,			collections.abc.Iterable						) else (stride, stride)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						padding if isinstance(_UpperCAmelCase,			collections.abc.Iterable						) else (padding, padding)

														SCREAMING_SNAKE_CASE__       :      str     =						nn.Convad(_UpperCAmelCase,			_UpperCAmelCase,			kernel_size=_UpperCAmelCase,			stride=_UpperCAmelCase,			padding=_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						norm_layer(_UpperCAmelCase						) if norm_layer else nn.Identity()






							def  A_     ( self     :					Any,			_UpperCAmelCase     :					Tuple						)				->			Union[str, Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.projection(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      str     =						self.norm(_UpperCAmelCase						)
														return embeddings


class 						lowerCamelCase       (nn.GroupNorm      ):







							"""simple docstring"""




							def __init__( self     :					str,			_UpperCAmelCase     :					Tuple,			**_UpperCAmelCase     :					Tuple						)				->			Union[str, Any]:





														"""simple docstring"""
														super().__init__(1,			_UpperCAmelCase,			**_UpperCAmelCase						)


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					Tuple,			_UpperCAmelCase     :					Optional[Any]						)				->			List[str]:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      str     =						nn.AvgPoolad(_UpperCAmelCase,			stride=1,			padding=pool_size // 2,			count_include_pad=_UpperCAmelCase						)






							def  A_     ( self     :					List[str],			_UpperCAmelCase     :					Union[str, Any]						)				->			int:





														"""simple docstring"""
														return self.pool(_UpperCAmelCase						) - hidden_states


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					str,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Union[str, Any]						)				->			Any:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      Any     =						nn.Convad(_UpperCAmelCase,			_UpperCAmelCase,			1						)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						nn.Convad(_UpperCAmelCase,			_UpperCAmelCase,			1						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						PoolFormerDropPath(_UpperCAmelCase						)
														if isinstance(config.hidden_act,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						ACTaFN[config.hidden_act]
														else:
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						config.hidden_act






							def  A_     ( self     :					Optional[Any],			_UpperCAmelCase     :					Optional[Any]						)				->			Optional[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						self.conva(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						self.act_fn(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.drop(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      int     =						self.conva(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.drop(_UpperCAmelCase						)

														return hidden_states


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					Any,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					List[str]						)				->			List[str]:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      int     =						PoolFormerPooling(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						PoolFormerOutput(_UpperCAmelCase,			_UpperCAmelCase,			_UpperCAmelCase,			_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						PoolFormerGroupNorm(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						PoolFormerGroupNorm(_UpperCAmelCase						)

														# Useful for training neural nets
														SCREAMING_SNAKE_CASE__       :      Tuple     =						PoolFormerDropPath(_UpperCAmelCase						) if drop_path > 0.0 else nn.Identity()
														SCREAMING_SNAKE_CASE__       :      Tuple     =						config.use_layer_scale
														if config.use_layer_scale:
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						nn.Parameter(
																					    config.layer_scale_init_value * torch.ones((_UpperCAmelCase)						),			requires_grad=_UpperCAmelCase						)
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						nn.Parameter(
																					    config.layer_scale_init_value * torch.ones((_UpperCAmelCase)						),			requires_grad=_UpperCAmelCase						)






							def  A_     ( self     :					Tuple,			_UpperCAmelCase     :					Any						)				->			Optional[Any]:





														"""simple docstring"""
														if self.use_layer_scale:
																					SCREAMING_SNAKE_CASE__       :      Any     =						self.pooling(self.before_norm(_UpperCAmelCase						)						)
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						self.layer_scale_a.unsqueeze(-1						).unsqueeze(-1						) * pooling_output
																					# First residual connection
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						hidden_states + self.drop_path(_UpperCAmelCase						)
																					SCREAMING_SNAKE_CASE__       :      int     =						()

																					SCREAMING_SNAKE_CASE__       :      str     =						self.output(self.after_norm(_UpperCAmelCase						)						)
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						self.layer_scale_a.unsqueeze(-1						).unsqueeze(-1						) * layer_output
																					# Second residual connection
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						hidden_states + self.drop_path(_UpperCAmelCase						)

																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						(output,) + outputs
																					return outputs

														else:
																					SCREAMING_SNAKE_CASE__       :      Dict     =						self.drop_path(self.pooling(self.before_norm(_UpperCAmelCase						)						)						)
																					# First residual connection
																					SCREAMING_SNAKE_CASE__       :      Any     =						pooling_output + hidden_states
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						()

																					# Second residual connection inside the PoolFormerOutput block
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						self.drop_path(self.output(self.after_norm(_UpperCAmelCase						)						)						)
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						hidden_states + layer_output

																					SCREAMING_SNAKE_CASE__       :      Tuple     =						(output,) + outputs
																					return outputs


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					Tuple,			_UpperCAmelCase     :					Dict						)				->			Union[str, Any]:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						config
														# stochastic depth decay rule
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						[x.item() for x in torch.linspace(0,			config.drop_path_rate,			sum(config.depths						)						)]

														# patch embeddings
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						[]
														for i in range(config.num_encoder_blocks						):
																					embeddings.append(
																					    PoolFormerEmbeddings(
																					        patch_size=config.patch_sizes[i],			stride=config.strides[i],			padding=config.padding[i],			num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1],			hidden_size=config.hidden_sizes[i],			)						)
														SCREAMING_SNAKE_CASE__       :      Dict     =						nn.ModuleList(_UpperCAmelCase						)

														# Transformer blocks
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						[]
														SCREAMING_SNAKE_CASE__       :      Any     =						0
														for i in range(config.num_encoder_blocks						):
																					# each block consists of layers
																					SCREAMING_SNAKE_CASE__       :      Dict     =						[]
																					if i != 0:
																												cur += config.depths[i - 1]
																					for j in range(config.depths[i]						):
																												layers.append(
																												    PoolFormerLayer(
																												        _UpperCAmelCase,			num_channels=config.hidden_sizes[i],			pool_size=config.pool_size,			hidden_size=config.hidden_sizes[i],			intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio						),			drop_path=dpr[cur + j],			)						)
																					blocks.append(nn.ModuleList(_UpperCAmelCase						)						)

														SCREAMING_SNAKE_CASE__       :      Any     =						nn.ModuleList(_UpperCAmelCase						)






							def  A_     ( self     :					Any,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Optional[Any]=False,			_UpperCAmelCase     :					List[str]=True						)				->			Union[str, Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						() if output_hidden_states else None

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						pixel_values
														for idx, layers in enumerate(zip(self.patch_embeddings,			self.block						)						):
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						layers
																					# Get patch embeddings from hidden_states
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						embedding_layer(_UpperCAmelCase						)
																					# Send the embeddings through the blocks
																					for _, blk in enumerate(_UpperCAmelCase						):
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						blk(_UpperCAmelCase						)
																												SCREAMING_SNAKE_CASE__       :      str     =						layer_outputs[0]

																					if output_hidden_states:
																												SCREAMING_SNAKE_CASE__       :      List[Any]     =						all_hidden_states + (hidden_states,)

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

														return BaseModelOutputWithNoAttention(last_hidden_state=_UpperCAmelCase,			hidden_states=_UpperCAmelCase						)


class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   PoolFormerConfig
							UpperCAmelCase_      						=   "poolformer"
							UpperCAmelCase_      						=   "pixel_values"
							UpperCAmelCase_      						=   True






							def  A_     ( self     :					Any,			_UpperCAmelCase     :					str						)				->			Dict:





														"""simple docstring"""
														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.LayerNorm						):
																					module.bias.data.zero_()
																					module.weight.data.fill_(1.0						)






							def  A_     ( self     :					Dict,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Dict=False						)				->			str:





														"""simple docstring"""
														if isinstance(_UpperCAmelCase,			_UpperCAmelCase						):
																					SCREAMING_SNAKE_CASE__       :      int     =						value


_lowerCamelCase :    List[str]  							= r'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): 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.
'''

_lowerCamelCase :    Tuple  							= r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
'''


@add_start_docstrings(
    "The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."							,     __lowerCamelCase							,     )
class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							def __init__( self     :					str,			_UpperCAmelCase     :					List[str]						)				->			Any:





														"""simple docstring"""
														super().__init__(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						config

														SCREAMING_SNAKE_CASE__       :      int     =						PoolFormerEncoder(_UpperCAmelCase						)

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






							def  A_     ( self     :					List[Any]						)				->			Optional[Any]:





														"""simple docstring"""
														return self.embeddings.patch_embeddings






							@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  A_     ( self     :					Dict,			_UpperCAmelCase     :					Optional[torch.FloatTensor] = None,			_UpperCAmelCase     :					Optional[bool] = None,			_UpperCAmelCase     :					Optional[bool] = None,			)				->			Union[Tuple, BaseModelOutputWithNoAttention]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						(
														    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
														)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						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"						)

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.encoder(
														    _UpperCAmelCase,			output_hidden_states=_UpperCAmelCase,			return_dict=_UpperCAmelCase,			)
														SCREAMING_SNAKE_CASE__       :      int     =						encoder_outputs[0]

														if not return_dict:
																					return (sequence_output, None) + encoder_outputs[1:]

														return BaseModelOutputWithNoAttention(
														    last_hidden_state=_UpperCAmelCase,			hidden_states=encoder_outputs.hidden_states,			)


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							def __init__( self     :					int,			_UpperCAmelCase     :					Tuple						)				->			Optional[int]:





														"""simple docstring"""
														super().__init__()
														SCREAMING_SNAKE_CASE__       :      Tuple     =						nn.Linear(config.hidden_size,			config.hidden_size						)






							def  A_     ( self     :					List[Any],			_UpperCAmelCase     :					int						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.dense(_UpperCAmelCase						)
														return output


@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "							,     __lowerCamelCase							,     )
class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							def __init__( self     :					str,			_UpperCAmelCase     :					str						)				->			List[str]:





														"""simple docstring"""
														super().__init__(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						config.num_labels
														SCREAMING_SNAKE_CASE__       :      Any     =						PoolFormerModel(_UpperCAmelCase						)

														# Final norm
														SCREAMING_SNAKE_CASE__       :      str     =						PoolFormerGroupNorm(config.hidden_sizes[-1]						)
														# Classifier head
														SCREAMING_SNAKE_CASE__       :      int     =						(
														    nn.Linear(config.hidden_sizes[-1],			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  A_     ( self     :					Optional[int],			_UpperCAmelCase     :					Optional[torch.FloatTensor] = None,			_UpperCAmelCase     :					Optional[torch.LongTensor] = None,			_UpperCAmelCase     :					Optional[bool] = None,			_UpperCAmelCase     :					Optional[bool] = None,			)				->			Union[Tuple, ImageClassifierOutputWithNoAttention]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						return_dict if return_dict is not None else self.config.use_return_dict

														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.poolformer(
														    _UpperCAmelCase,			output_hidden_states=_UpperCAmelCase,			return_dict=_UpperCAmelCase,			)

														SCREAMING_SNAKE_CASE__       :      Dict     =						outputs[0]

														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						self.classifier(self.norm(_UpperCAmelCase						).mean([-2, -1]						)						)

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

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

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

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













import collections
import gzip
import os
import urllib

import numpy
from tensorflow.python.framework import dtypes, random_seed
from tensorflow.python.platform import gfile
from tensorflow.python.util.deprecation import deprecated

_lowerCamelCase :    Tuple  							= collections.namedtuple('''_Datasets''', ['''train''', '''validation''', '''test'''])

# CVDF mirror of http://yann.lecun.com/exdb/mnist/
_lowerCamelCase :    Optional[Any]  							= '''https://storage.googleapis.com/cvdf-datasets/mnist/'''




def _a (  SCREAMING_SNAKE_CASE__      :  Any       )	->    List[Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						numpy.dtype(numpy.uintaa       ).newbyteorder(">"       )
							return numpy.frombuffer(bytestream.read(4       )			,				dtype=SCREAMING_SNAKE_CASE__       )[0]




@deprecated(SCREAMING_SNAKE_CASE__			,				"Please use tf.data to implement this functionality."       )
def _a (  SCREAMING_SNAKE_CASE__      :  Optional[Any]       )	->    List[str]:

							'''simple docstring'''






							print("Extracting"			,				f.name       )
							with gzip.GzipFile(fileobj=SCREAMING_SNAKE_CASE__       ) as bytestream:
														SCREAMING_SNAKE_CASE__       :      int     =						_readaa(SCREAMING_SNAKE_CASE__       )
														if magic != 20_51:
																					raise ValueError(
																					    "Invalid magic number %d in MNIST image file: %s" % (magic, f.name)       )
														SCREAMING_SNAKE_CASE__       :      Dict     =						_readaa(SCREAMING_SNAKE_CASE__       )
														SCREAMING_SNAKE_CASE__       :      str     =						_readaa(SCREAMING_SNAKE_CASE__       )
														SCREAMING_SNAKE_CASE__       :      Dict     =						_readaa(SCREAMING_SNAKE_CASE__       )
														SCREAMING_SNAKE_CASE__       :      str     =						bytestream.read(rows * cols * num_images       )
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						numpy.frombuffer(SCREAMING_SNAKE_CASE__			,				dtype=numpy.uinta       )
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						data.reshape(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				1       )
														return data




@deprecated(SCREAMING_SNAKE_CASE__			,				"Please use tf.one_hot on tensors."       )
def _a (  SCREAMING_SNAKE_CASE__      :  str			,				SCREAMING_SNAKE_CASE__      :  Tuple       )	->    Union[str, Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						labels_dense.shape[0]
							SCREAMING_SNAKE_CASE__       :      List[str]     =						numpy.arange(SCREAMING_SNAKE_CASE__       ) * num_classes
							SCREAMING_SNAKE_CASE__       :      Dict     =						numpy.zeros((num_labels, num_classes)       )
							SCREAMING_SNAKE_CASE__       :      Tuple     =						1
							return labels_one_hot




@deprecated(SCREAMING_SNAKE_CASE__			,				"Please use tf.data to implement this functionality."       )
def _a (  SCREAMING_SNAKE_CASE__      :  List[Any]			,				SCREAMING_SNAKE_CASE__      :  Union[str, Any]=False			,				SCREAMING_SNAKE_CASE__      :  Optional[int]=10       )	->    Optional[int]:

							'''simple docstring'''






							print("Extracting"			,				f.name       )
							with gzip.GzipFile(fileobj=SCREAMING_SNAKE_CASE__       ) as bytestream:
														SCREAMING_SNAKE_CASE__       :      List[str]     =						_readaa(SCREAMING_SNAKE_CASE__       )
														if magic != 20_49:
																					raise ValueError(
																					    "Invalid magic number %d in MNIST label file: %s" % (magic, f.name)       )
														SCREAMING_SNAKE_CASE__       :      Tuple     =						_readaa(SCREAMING_SNAKE_CASE__       )
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						bytestream.read(SCREAMING_SNAKE_CASE__       )
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						numpy.frombuffer(SCREAMING_SNAKE_CASE__			,				dtype=numpy.uinta       )
														if one_hot:
																					return _dense_to_one_hot(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
														return labels


class 						lowerCamelCase       :







							"""simple docstring"""




							@deprecated(
							    _UpperCAmelCase,			"Please use alternatives such as official/mnist/_DataSet.py"
							    " from tensorflow/models.",			)
							def __init__( self     :					int,			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					List[str]=False,			_UpperCAmelCase     :					Tuple=False,			_UpperCAmelCase     :					Union[str, Any]=dtypes.floataa,			_UpperCAmelCase     :					Optional[Any]=True,			_UpperCAmelCase     :					Optional[int]=None,			)				->			List[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						random_seed.get_seed(_UpperCAmelCase						)
														# If op level seed is not set, use whatever graph level seed is returned
														numpy.random.seed(seeda if seed is None else seeda						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						dtypes.as_dtype(_UpperCAmelCase						).base_dtype
														if dtype not in (dtypes.uinta, dtypes.floataa):
																					raise TypeError("Invalid image dtype %r, expected uint8 or float32" % dtype						)
														if fake_data:
																					SCREAMING_SNAKE_CASE__       :      str     =						1_0_0_0_0
																					SCREAMING_SNAKE_CASE__       :      str     =						one_hot
														else:
																					assert (
																					    images.shape[0] == labels.shape[0]
																					), F'''images.shape: {images.shape} labels.shape: {labels.shape}'''
																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						images.shape[0]

																					# Convert shape from [num examples, rows, columns, depth]
																					# to [num examples, rows*columns] (assuming depth == 1)
																					if reshape:
																												assert images.shape[3] == 1
																												SCREAMING_SNAKE_CASE__       :      List[Any]     =						images.reshape(
																												    images.shape[0],			images.shape[1] * images.shape[2]						)
																					if dtype == dtypes.floataa:
																												# Convert from [0, 255] -> [0.0, 1.0].
																												SCREAMING_SNAKE_CASE__       :      Any     =						images.astype(numpy.floataa						)
																												SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						numpy.multiply(_UpperCAmelCase,			1.0 / 255.0						)
														SCREAMING_SNAKE_CASE__       :      List[str]     =						images
														SCREAMING_SNAKE_CASE__       :      str     =						labels
														SCREAMING_SNAKE_CASE__       :      Any     =						0
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						0






							@property
							def  A_     ( self     :					List[Any]						)				->			int:





														"""simple docstring"""
														return self._images






							@property
							def  A_     ( self     :					List[str]						)				->			Optional[int]:





														"""simple docstring"""
														return self._labels






							@property
							def  A_     ( self     :					Optional[Any]						)				->			Union[str, Any]:





														"""simple docstring"""
														return self._num_examples






							@property
							def  A_     ( self     :					List[Any]						)				->			Tuple:





														"""simple docstring"""
														return self._epochs_completed






							def  A_     ( self     :					List[str],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					List[Any]=False,			_UpperCAmelCase     :					str=True						)				->			Any:





														"""simple docstring"""
														if fake_data:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						[1] * 7_8_4
																					SCREAMING_SNAKE_CASE__       :      Dict     =						[1] + [0] * 9 if self.one_hot else 0
																					return (
																					    [fake_image for _ in range(_UpperCAmelCase						)],
																					    [fake_label for _ in range(_UpperCAmelCase						)],
																					)
														SCREAMING_SNAKE_CASE__       :      str     =						self._index_in_epoch
														# Shuffle for the first epoch
														if self._epochs_completed == 0 and start == 0 and shuffle:
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						numpy.arange(self._num_examples						)
																					numpy.random.shuffle(_UpperCAmelCase						)
																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.images[perma]
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.labels[perma]
														# Go to the next epoch
														if start + batch_size > self._num_examples:
																					# Finished epoch
																					self._epochs_completed += 1
																					# Get the rest examples in this epoch
																					SCREAMING_SNAKE_CASE__       :      Dict     =						self._num_examples - start
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						self._images[start : self._num_examples]
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self._labels[start : self._num_examples]
																					# Shuffle the data
																					if shuffle:
																												SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						numpy.arange(self._num_examples						)
																												numpy.random.shuffle(_UpperCAmelCase						)
																												SCREAMING_SNAKE_CASE__       :      Dict     =						self.images[perm]
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						self.labels[perm]
																					# Start next epoch
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						0
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						batch_size - rest_num_examples
																					SCREAMING_SNAKE_CASE__       :      Dict     =						self._index_in_epoch
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self._images[start:end]
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self._labels[start:end]
																					return (
																					    numpy.concatenate((images_rest_part, images_new_part),			axis=0						),
																					    numpy.concatenate((labels_rest_part, labels_new_part),			axis=0						),
																					)
														else:
																					self._index_in_epoch += batch_size
																					SCREAMING_SNAKE_CASE__       :      Tuple     =						self._index_in_epoch
																					return self._images[start:end], self._labels[start:end]




@deprecated(SCREAMING_SNAKE_CASE__			,				"Please write your own downloading logic."       )
def _a (  SCREAMING_SNAKE_CASE__      :  Union[str, Any]			,				SCREAMING_SNAKE_CASE__      :  str			,				SCREAMING_SNAKE_CASE__      :  Union[str, Any]       )	->    str:

							'''simple docstring'''






							if not gfile.Exists(SCREAMING_SNAKE_CASE__       ):
														gfile.MakeDirs(SCREAMING_SNAKE_CASE__       )
							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						os.path.join(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )
							if not gfile.Exists(SCREAMING_SNAKE_CASE__       ):
														urllib.request.urlretrieve(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__       )  # noqa: S310
														with gfile.GFile(SCREAMING_SNAKE_CASE__       ) as f:
																					SCREAMING_SNAKE_CASE__       :      Any     =						f.size()
														print("Successfully downloaded"			,				SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				"bytes."       )
							return filepath







@deprecated(
    SCREAMING_SNAKE_CASE__			,				"Please use alternatives such as:" " tensorflow_datasets.load('mnist')"       )
def _a (  SCREAMING_SNAKE_CASE__      :  List[str]			,				SCREAMING_SNAKE_CASE__      :  str=False			,				SCREAMING_SNAKE_CASE__      :  Any=False			,				SCREAMING_SNAKE_CASE__      :  List[Any]=dtypes.floataa			,				SCREAMING_SNAKE_CASE__      :  str=True			,				SCREAMING_SNAKE_CASE__      :  Dict=50_00			,				SCREAMING_SNAKE_CASE__      :  str=None			,				SCREAMING_SNAKE_CASE__      :  Optional[int]=DEFAULT_SOURCE_URL			,				)	->    Any:

							'''simple docstring'''






							if fake_data:

														def fake():
																					return _DataSet(
																					    []			,				[]			,				fake_data=SCREAMING_SNAKE_CASE__			,				one_hot=SCREAMING_SNAKE_CASE__			,				dtype=SCREAMING_SNAKE_CASE__			,				seed=SCREAMING_SNAKE_CASE__       )

														SCREAMING_SNAKE_CASE__       :      int     =						fake()
														SCREAMING_SNAKE_CASE__       :      List[str]     =						fake()
														SCREAMING_SNAKE_CASE__       :      Any     =						fake()
														return _Datasets(train=SCREAMING_SNAKE_CASE__			,				validation=SCREAMING_SNAKE_CASE__			,				test=SCREAMING_SNAKE_CASE__       )

							if not source_url:  # empty string check
														SCREAMING_SNAKE_CASE__       :      str     =						DEFAULT_SOURCE_URL

							SCREAMING_SNAKE_CASE__       :      Dict     =						"train-images-idx3-ubyte.gz"
							SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						"train-labels-idx1-ubyte.gz"
							SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						"t10k-images-idx3-ubyte.gz"
							SCREAMING_SNAKE_CASE__       :      str     =						"t10k-labels-idx1-ubyte.gz"

							SCREAMING_SNAKE_CASE__       :      List[Any]     =						_maybe_download(
							    SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				source_url + train_images_file       )
							with gfile.Open(SCREAMING_SNAKE_CASE__			,				"rb"       ) as f:
														SCREAMING_SNAKE_CASE__       :      Dict     =						_extract_images(SCREAMING_SNAKE_CASE__       )

							SCREAMING_SNAKE_CASE__       :      List[Any]     =						_maybe_download(
							    SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				source_url + train_labels_file       )
							with gfile.Open(SCREAMING_SNAKE_CASE__			,				"rb"       ) as f:
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						_extract_labels(SCREAMING_SNAKE_CASE__			,				one_hot=SCREAMING_SNAKE_CASE__       )

							SCREAMING_SNAKE_CASE__       :      Dict     =						_maybe_download(
							    SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				source_url + test_images_file       )
							with gfile.Open(SCREAMING_SNAKE_CASE__			,				"rb"       ) as f:
														SCREAMING_SNAKE_CASE__       :      List[str]     =						_extract_images(SCREAMING_SNAKE_CASE__       )

							SCREAMING_SNAKE_CASE__       :      Optional[int]     =						_maybe_download(
							    SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				source_url + test_labels_file       )
							with gfile.Open(SCREAMING_SNAKE_CASE__			,				"rb"       ) as f:
														SCREAMING_SNAKE_CASE__       :      Any     =						_extract_labels(SCREAMING_SNAKE_CASE__			,				one_hot=SCREAMING_SNAKE_CASE__       )

							if not 0 <= validation_size <= len(SCREAMING_SNAKE_CASE__       ):
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						(
														    "Validation size should be between 0 and "
														    f'''{len(SCREAMING_SNAKE_CASE__       )}. Received: {validation_size}.'''
														)
														raise ValueError(SCREAMING_SNAKE_CASE__       )

							SCREAMING_SNAKE_CASE__       :      str     =						train_images[:validation_size]
							SCREAMING_SNAKE_CASE__       :      str     =						train_labels[:validation_size]
							SCREAMING_SNAKE_CASE__       :      Any     =						train_images[validation_size:]
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						train_labels[validation_size:]

							SCREAMING_SNAKE_CASE__       :      Any     =						{"dtype": dtype, "reshape": reshape, "seed": seed}

							SCREAMING_SNAKE_CASE__       :      Tuple     =						_DataSet(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				**SCREAMING_SNAKE_CASE__       )
							SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						_DataSet(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				**SCREAMING_SNAKE_CASE__       )
							SCREAMING_SNAKE_CASE__       :      str     =						_DataSet(SCREAMING_SNAKE_CASE__			,				SCREAMING_SNAKE_CASE__			,				**SCREAMING_SNAKE_CASE__       )

							return _Datasets(train=SCREAMING_SNAKE_CASE__			,				validation=SCREAMING_SNAKE_CASE__			,				test=SCREAMING_SNAKE_CASE__       )






