



"""simple docstring"""

from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
     import PIL


lowercase__					:						int												=			logging.get_logger(__name__)




class  _UpperCAmelCase (    lowerCAmelCase__):
 _lowerCAmelCase     :    Union[str, Any]	      =	["""pixel_values"""]

 def __init__(  self	: Tuple						,  lowercase_	: bool = True						,  lowercase_	: Dict[str, int] = None						,  lowercase_	: float = None						,  lowercase_	: PILImageResampling = PILImageResampling.BILINEAR						,  lowercase_	: bool = True						,  lowercase_	: Union[int, float] = 1 / 255						,  lowercase_	: bool = True						,  lowercase_	: Optional[Union[float, List[float]]] = None						,  lowercase_	: Optional[Union[float, List[float]]] = None						,  **lowercase_	: Optional[Any]						,  ):
     super().__init__(**lowercase_       )
     snake_case_				:  Any       =			size if size is not None else {'''shortest_edge''': 384}
     snake_case_				:  Any       =			get_size_dict(lowercase_						,  default_to_square=lowercase_       )

     snake_case_				:  Union[str, Any]       =			do_resize
     snake_case_				:  int       =			size
     # Default value set here for backwards compatibility where the value in config is None
     snake_case_				:  str       =			crop_pct if crop_pct is not None else 224 / 256
     snake_case_				:  Optional[int]       =			resample
     snake_case_				:  Dict       =			do_rescale
     snake_case_				:  int       =			rescale_factor
     snake_case_				:  Dict       =			do_normalize
     snake_case_				:  Dict       =			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
     snake_case_				:  List[Any]       =			image_std if image_std is not None else IMAGENET_STANDARD_STD

 def 		_snake_case      (  self	: Any						,  lowercase_	: np.ndarray						,  lowercase_	: Dict[str, int]						,  lowercase_	: float						,  lowercase_	: PILImageResampling = PILImageResampling.BICUBIC						,  lowercase_	: Optional[Union[str, ChannelDimension]] = None						,  **lowercase_	: List[str]						,  ):
     snake_case_				:  Tuple       =			get_size_dict(lowercase_						,  default_to_square=lowercase_       )
     if "shortest_edge" not in size:
         raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"       )
     snake_case_				:  Union[str, Any]       =			size['''shortest_edge''']

     if shortest_edge < 384:
         # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
         snake_case_				:  Tuple       =			int(shortest_edge / crop_pct       )
         snake_case_				:  List[str]       =			get_resize_output_image_size(lowercase_						,  size=lowercase_						,  default_to_square=lowercase_       )
         snake_case_				:  List[Any]       =			resize(image=lowercase_						,  size=lowercase_						,  resample=lowercase_						,  data_format=lowercase_						,  **lowercase_       )
         # then crop to (shortest_edge, shortest_edge)
         return center_crop(image=lowercase_						,  size=(shortest_edge, shortest_edge)						,  data_format=lowercase_						,  **lowercase_       )
     else:
         # warping (no cropping) when evaluated at 384 or larger
         return resize(
             lowercase_						,  size=(shortest_edge, shortest_edge)						,  resample=lowercase_						,  data_format=lowercase_						,  **lowercase_       )

 def 		_snake_case      (  self	: int						,  lowercase_	: np.ndarray						,  lowercase_	: Union[int, float]						,  lowercase_	: Optional[Union[str, ChannelDimension]] = None						,  **lowercase_	: List[Any]						,  ):
     return rescale(lowercase_						,  scale=lowercase_						,  data_format=lowercase_						,  **lowercase_       )

 def 		_snake_case      (  self	: Union[str, Any]						,  lowercase_	: np.ndarray						,  lowercase_	: Union[float, List[float]]						,  lowercase_	: Union[float, List[float]]						,  lowercase_	: Optional[Union[str, ChannelDimension]] = None						,  **lowercase_	: Tuple						,  ):
     return normalize(lowercase_						,  mean=lowercase_						,  std=lowercase_						,  data_format=lowercase_						,  **lowercase_       )






 def 		_snake_case      (  self	: Any						,  lowercase_	: ImageInput						,  lowercase_	: bool = None						,  lowercase_	: Dict[str, int] = None						,  lowercase_	: float = None						,  lowercase_	: PILImageResampling = None						,  lowercase_	: bool = None						,  lowercase_	: float = None						,  lowercase_	: bool = None						,  lowercase_	: Optional[Union[float, List[float]]] = None						,  lowercase_	: Optional[Union[float, List[float]]] = None						,  lowercase_	: Optional[Union[str, TensorType]] = None						,  lowercase_	: ChannelDimension = ChannelDimension.FIRST						,  **lowercase_	: Any						,  ):
     snake_case_				:  Any       =			do_resize if do_resize is not None else self.do_resize
     snake_case_				:  List[Any]       =			crop_pct if crop_pct is not None else self.crop_pct
     snake_case_				:  Dict       =			resample if resample is not None else self.resample
     snake_case_				:  List[str]       =			do_rescale if do_rescale is not None else self.do_rescale
     snake_case_				:  Union[str, Any]       =			rescale_factor if rescale_factor is not None else self.rescale_factor
     snake_case_				:  Dict       =			do_normalize if do_normalize is not None else self.do_normalize
     snake_case_				:  List[Any]       =			image_mean if image_mean is not None else self.image_mean
     snake_case_				:  Optional[Any]       =			image_std if image_std is not None else self.image_std

     snake_case_				:  int       =			size if size is not None else self.size
     snake_case_				:  Dict       =			get_size_dict(lowercase_						,  default_to_square=lowercase_       )

     snake_case_				:  List[Any]       =			make_list_of_images(lowercase_       )

     if not valid_images(lowercase_       ):
         raise ValueError(
             '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
             '''torch.Tensor, tf.Tensor or jax.ndarray.'''       )

     if do_resize and size is None or resample is None:
         raise ValueError('''Size and resample must be specified if do_resize is True.'''       )

     if do_resize and size["shortest_edge"] < 384 and crop_pct is None:
         raise ValueError('''crop_pct must be specified if size < 384.'''       )

     if do_rescale and rescale_factor is None:
         raise ValueError('''Rescale factor must be specified if do_rescale is True.'''       )

     if do_normalize and (image_mean is None or image_std is None):
         raise ValueError('''Image mean and std must be specified if do_normalize is True.'''       )

     # All transformations expect numpy arrays.
     snake_case_				:  Dict       =			[to_numpy_array(lowercase_       ) for image in images]

     if do_resize:
         snake_case_				:  Optional[Any]       =			[self.resize(image=lowercase_						,  size=lowercase_						,  crop_pct=lowercase_						,  resample=lowercase_       ) for image in images]

     if do_rescale:
         snake_case_				:  int       =			[self.rescale(image=lowercase_						,  scale=lowercase_       ) for image in images]

     if do_normalize:
         snake_case_				:  Optional[Any]       =			[self.normalize(image=lowercase_						,  mean=lowercase_						,  std=lowercase_       ) for image in images]

     snake_case_				:  Optional[Any]       =			[to_channel_dimension_format(lowercase_						,  lowercase_       ) for image in images]

     snake_case_				:  List[str]       =			{'''pixel_values''': images}
     return BatchFeature(data=lowercase_						,  tensor_type=lowercase_       )










"""simple docstring"""

import tensorflow as tf

from ...tf_utils import shape_list




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

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

     snake_case_				:  int       =			vocab_size
     snake_case_				:  Union[str, Any]       =			d_embed
     snake_case_				:  Optional[int]       =			d_proj

     snake_case_				:  int       =			cutoffs + [vocab_size]
     snake_case_				:  Optional[int]       =			[0] + self.cutoffs
     snake_case_				:  List[str]       =			div_val

     snake_case_				:  int       =			self.cutoffs[0]
     snake_case_				:  Optional[Any]       =			len(self.cutoffs       ) - 1
     snake_case_				:  Any       =			self.shortlist_size + self.n_clusters
     snake_case_				:  Dict       =			keep_order

     snake_case_				:  Tuple       =			[]
     snake_case_				:  Optional[Any]       =			[]

 def 		_snake_case      (  self	: Dict						,  lowercase_	: int       ):
     if self.n_clusters > 0:
         snake_case_				:  Tuple       =			self.add_weight(
             shape=(self.n_clusters, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name='''cluster_weight'''       )
         snake_case_				:  List[str]       =			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:
                 snake_case_				:  List[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_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(self.vocab_size, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  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       )       ):
             snake_case_,			snake_case_				:  str       =			self.cutoff_ends[i], self.cutoff_ends[i + 1]
             snake_case_				:  List[str]       =			self.d_embed // (self.div_val**i)

             snake_case_				:  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_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(r_idx - l_idx, d_emb_i)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  List[Any]       =			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 		_snake_case      (  lowercase_	: Dict						,  lowercase_	: Union[str, Any]						,  lowercase_	: Optional[int]						,  lowercase_	: str=None       ):
     snake_case_				:  str       =			x
     if proj is not None:
         snake_case_				:  str       =			tf.einsum('''ibd,ed->ibe'''						,  lowercase_						,  lowercase_       )
     return tf.einsum('''ibd,nd->ibn'''						,  lowercase_						,  lowercase_       ) + b

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






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

             if self.div_val == 1:
                 snake_case_				:  Optional[int]       =			self.out_layers[0][0][l_idx:r_idx]
                 snake_case_				:  List[Any]       =			self.out_layers[0][1][l_idx:r_idx]
             else:
                 snake_case_				:  Dict       =			self.out_layers[i][0]
                 snake_case_				:  str       =			self.out_layers[i][1]

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

                 snake_case_				:  List[Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[0]       )
                 snake_case_				:  Union[str, Any]       =			tf.nn.log_softmax(lowercase_       )
                 out.append(head_logprob[..., : self.cutoffs[0]]       )
                 if target is not None:
                     snake_case_				:  int       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  str       =			self._gather_logprob(lowercase_						,  lowercase_       )
             else:
                 snake_case_				:  Union[str, Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[i]       )
                 snake_case_				:  List[str]       =			tf.nn.log_softmax(lowercase_       )
                 snake_case_				:  Dict       =			self.cutoffs[0] + i - 1  # No probability for the head cluster
                 snake_case_				:  Union[str, Any]       =			head_logprob[..., cluster_prob_idx, None] + tail_logprob
                 out.append(lowercase_       )
                 if target is not None:
                     snake_case_				:  Optional[int]       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Dict       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Union[str, Any]       =			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_       )       )
         snake_case_				:  Union[str, Any]       =			tf.concat(lowercase_						,  axis=-1       )

     if target is not None:
         if return_mean:
             snake_case_				:  Dict       =			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





