

"""simple docstring"""





import tensorflow as tf

from ...tf_utils import shape_list







class 				SCREAMING_SNAKE_CASE					(		tf.keras.layers.Layer       ):






					"""simple docstring"""



					def __init__(       self			:				List[Any]   ,lowercase_			:				Tuple   ,lowercase_			:				List[Any]   ,lowercase_			:				List[str]   ,lowercase_			:				Optional[int]   ,lowercase_			:				str=1   ,lowercase_			:				Union[str, Any]=False   ,**lowercase_			:				Union[str, Any]		):
						super().__init__(**lowercase_		)

						lowerCAmelCase__	:						Any  				=	vocab_size
						lowerCAmelCase__	:						Tuple  				=	d_embed
						lowerCAmelCase__	:						Union[str, Any]  				=	d_proj

						lowerCAmelCase__	:						Union[str, Any]  				=	cutoffs + [vocab_size]
						lowerCAmelCase__	:						Optional[Any]  				=	[0] + self.cutoffs
						lowerCAmelCase__	:						int  				=	div_val

						lowerCAmelCase__	:						Optional[Any]  				=	self.cutoffs[0]
						lowerCAmelCase__	:						Tuple  				=	len(self.cutoffs		) - 1
						lowerCAmelCase__	:						Any  				=	self.shortlist_size + self.n_clusters
						lowerCAmelCase__	:						str  				=	keep_order

						lowerCAmelCase__	:						List[str]  				=	[]
						lowerCAmelCase__	:						List[str]  				=	[]







					def        __lowerCAmelCase (       self			:				Optional[int]   ,lowercase_			:				str		):
						if self.n_clusters > 0:
							lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
							    shape=(self.n_clusters, self.d_embed)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name='''cluster_weight'''		)
							lowerCAmelCase__	:						List[Any]  				=	self.add_weight(
							    shape=(self.n_clusters,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name='''cluster_bias'''		)

						if self.div_val == 1:
							for i in range(len(self.cutoffs		)		):
								if self.d_proj != self.d_embed:
									lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
									    shape=(self.d_embed, self.d_proj)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_projs_._{i}'   ,)
									self.out_projs.append(lowercase_		)
								else:
									self.out_projs.append(lowercase_		)
								lowerCAmelCase__	:						Any  				=	self.add_weight(
								    shape=(self.vocab_size, self.d_embed)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._weight'   ,)
								lowerCAmelCase__	:						Union[str, Any]  				=	self.add_weight(
								    shape=(self.vocab_size,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._bias'   ,)
								self.out_layers.append((weight, bias)		)
						else:
							for i in range(len(self.cutoffs		)		):
								lowerCAmelCase__				,lowerCAmelCase__	:						Union[str, Any]  				=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
								lowerCAmelCase__	:						int  				=	self.d_embed // (self.div_val**i)

								lowerCAmelCase__	:						Optional[Any]  				=	self.add_weight(
								    shape=(d_emb_i, self.d_proj)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_projs_._{i}'		)
								self.out_projs.append(lowercase_		)
								lowerCAmelCase__	:						str  				=	self.add_weight(
								    shape=(r_idx - l_idx, d_emb_i)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._weight'   ,)
								lowerCAmelCase__	:						Optional[int]  				=	self.add_weight(
								    shape=(r_idx - l_idx,)   ,initializer='''zeros'''   ,trainable=lowercase_   ,name=F'out_layers_._{i}_._bias'   ,)
								self.out_layers.append((weight, bias)		)
						super().build(lowercase_		)







					@staticmethod
					def        __lowerCAmelCase (       lowercase_			:				List[str]   ,lowercase_			:				int   ,lowercase_			:				Tuple   ,lowercase_			:				List[Any]=None		):
						lowerCAmelCase__	:						Optional[Any]  				=	x
						if proj is not None:
							lowerCAmelCase__	:						Optional[Any]  				=	tf.einsum('''ibd,ed->ibe'''   ,lowercase_   ,lowercase_		)
						return tf.einsum('''ibd,nd->ibn'''   ,lowercase_   ,lowercase_		) + b







					@staticmethod
					def        __lowerCAmelCase (       lowercase_			:				int   ,lowercase_			:				Tuple		):
						lowerCAmelCase__	:						Optional[int]  				=	shape_list(lowercase_		)
						lowerCAmelCase__	:						Dict  				=	tf.range(lp_size[0]   ,dtype=target.dtype		)
						lowerCAmelCase__	:						str  				=	tf.stack([r, target]   ,1		)
						return tf.gather_nd(lowercase_   ,lowercase_		)







					def        __lowerCAmelCase (       self			:				Optional[int]   ,lowercase_			:				str   ,lowercase_			:				Dict   ,lowercase_			:				Dict=True   ,lowercase_			:				Dict=False		):
						lowerCAmelCase__	:						Optional[Any]  				=	0
						if self.n_clusters == 0:
							lowerCAmelCase__	:						Union[str, Any]  				=	self._logit(lowercase_   ,self.out_layers[0][0]   ,self.out_layers[0][1]   ,self.out_projs[0]		)
							if target is not None:
								lowerCAmelCase__	:						Any  				=	tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowercase_   ,logits=lowercase_		)
							lowerCAmelCase__	:						Any  				=	tf.nn.log_softmax(lowercase_   ,axis=-1		)
						else:
							lowerCAmelCase__	:						Optional[Any]  				=	shape_list(lowercase_		)
							lowerCAmelCase__	:						Union[str, Any]  				=	[]
							lowerCAmelCase__	:						int  				=	tf.zeros(hidden_sizes[:2]		)
							for i in range(len(self.cutoffs		)		):
								lowerCAmelCase__				,lowerCAmelCase__	:						Any  				=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
								if target is not None:
									lowerCAmelCase__	:						int  				=	(target >= l_idx) & (target < r_idx)
									lowerCAmelCase__	:						List[Any]  				=	tf.where(lowercase_		)
									lowerCAmelCase__	:						Union[str, Any]  				=	tf.boolean_mask(lowercase_   ,lowercase_		) - l_idx

								if self.div_val == 1:
									lowerCAmelCase__	:						int  				=	self.out_layers[0][0][l_idx:r_idx]
									lowerCAmelCase__	:						Union[str, Any]  				=	self.out_layers[0][1][l_idx:r_idx]
								else:
									lowerCAmelCase__	:						Optional[Any]  				=	self.out_layers[i][0]
									lowerCAmelCase__	:						Optional[int]  				=	self.out_layers[i][1]

								if i == 0:
									lowerCAmelCase__	:						Union[str, Any]  				=	tf.concat([cur_W, self.cluster_weight]   ,0		)
									lowerCAmelCase__	:						Tuple  				=	tf.concat([cur_b, self.cluster_bias]   ,0		)

									lowerCAmelCase__	:						List[Any]  				=	self._logit(lowercase_   ,lowercase_   ,lowercase_   ,self.out_projs[0]		)
									lowerCAmelCase__	:						List[str]  				=	tf.nn.log_softmax(lowercase_		)
									out.append(head_logprob[..., : self.cutoffs[0]]		)
									if target is not None:
										lowerCAmelCase__	:						List[str]  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						Union[str, Any]  				=	self._gather_logprob(lowercase_   ,lowercase_		)
								else:
									lowerCAmelCase__	:						int  				=	self._logit(lowercase_   ,lowercase_   ,lowercase_   ,self.out_projs[i]		)
									lowerCAmelCase__	:						Tuple  				=	tf.nn.log_softmax(lowercase_		)
									lowerCAmelCase__	:						Tuple  				=	self.cutoffs[0] + i - 1  # No probability for the head cluster
									lowerCAmelCase__	:						Optional[int]  				=	head_logprob[..., cluster_prob_idx, None] + tail_logprob
									out.append(lowercase_		)
									if target is not None:
										lowerCAmelCase__	:						Optional[Any]  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						Dict  				=	tf.boolean_mask(lowercase_   ,lowercase_		)
										lowerCAmelCase__	:						str  				=	self._gather_logprob(lowercase_   ,lowercase_		)
										cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
								if target is not None:
									loss += tf.scatter_nd(lowercase_   ,-cur_logprob   ,shape_list(lowercase_		)		)
							lowerCAmelCase__	:						Optional[int]  				=	tf.concat(lowercase_   ,axis=-1		)

						if target is not None:
							if return_mean:
								lowerCAmelCase__	:						List[Any]  				=	tf.reduce_mean(lowercase_		)
							# Add the training-time loss value to the layer using `self.add_loss()`.
							self.add_loss(lowercase_		)

							# Log the loss as a metric (we could log arbitrary metrics,
							# including different metrics for training and inference.
							self.add_metric(lowercase_   ,name=self.name   ,aggregation='''mean''' if return_mean else ''''''		)

						return out






import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse('9.1.0'):
		lowercase  :					int													=       {
		    'linear': PIL.Image.Resampling.BILINEAR,
		    'bilinear': PIL.Image.Resampling.BILINEAR,
		    'bicubic': PIL.Image.Resampling.BICUBIC,
		    'lanczos': PIL.Image.Resampling.LANCZOS,
		    'nearest': PIL.Image.Resampling.NEAREST,
		}
else:
		lowercase  :					Any													=       {
		    'linear': PIL.Image.LINEAR,
		    'bilinear': PIL.Image.BILINEAR,
		    'bicubic': PIL.Image.BICUBIC,
		    'lanczos': PIL.Image.LANCZOS,
		    'nearest': PIL.Image.NEAREST,
		}




def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			int)					->    Any:






							'''simple docstring'''





							__UpperCamelCase	:   Union[str, Any]			   =					(images / 2 + 0.5).clamp(0						,  1)
							__UpperCamelCase	:   int			   =					images.cpu().permute(0						,  2						,  3						,  1).float().numpy()
							__UpperCamelCase	:   Dict			   =					numpy_to_pil(_lowerCamelCase)
							return images




def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Optional[Any])					->    Optional[Any]:






							'''simple docstring'''





							if images.ndim == 3:
														__UpperCamelCase	:   Optional[int]			   =					images[None, ...]
							__UpperCamelCase	:   Dict			   =					(images * 255).round().astype("uint8")
							if images.shape[-1] == 1:
														# special case for grayscale (single channel) images
														__UpperCamelCase	:   int			   =					[Image.fromarray(image.squeeze()						,  mode="L") for image in images]
							else:
														__UpperCamelCase	:   Union[str, Any]			   =					[Image.fromarray(_lowerCamelCase) for image in images]

							return pil_images