
'''simple docstring'''


from typing import Dict, Iterable, 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_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


snake_case_ :  Optional[int]     =					logging.get_logger(__name__)


class 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   _snake_case		   =							['''pixel_values''']







   def __init__(			self ,   lowerCamelCase__ = True ,   lowerCamelCase__ = None ,   lowerCamelCase__ = PILImageResampling.BICUBIC ,   lowerCamelCase__ = True ,   lowerCamelCase__ = None ,   lowerCamelCase__ = True ,   lowerCamelCase__ = 1 / 2_5_5 ,   lowerCamelCase__ = True ,   lowerCamelCase__ = IMAGENET_DEFAULT_MEAN ,   lowerCamelCase__ = IMAGENET_DEFAULT_STD ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      super().__init__(**lowerCamelCase__   )
      UpperCamelCase       =  size if size is not None else {'''shortest_edge''': 2_2_4}
      UpperCamelCase       =  get_size_dict(lowerCamelCase__ ,   default_to_square=lowerCamelCase__   )
      UpperCamelCase       =  crop_size if crop_size is not None else {'''height''': 2_2_4, '''width''': 2_2_4}
      UpperCamelCase       =  get_size_dict(lowerCamelCase__ ,   param_name='''crop_size'''   )

      UpperCamelCase       =  do_resize
      UpperCamelCase       =  size
      UpperCamelCase       =  resample
      UpperCamelCase       =  do_center_crop
      UpperCamelCase       =  crop_size
      UpperCamelCase       =  do_rescale
      UpperCamelCase       =  rescale_factor
      UpperCamelCase       =  do_normalize
      UpperCamelCase       =  image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
      UpperCamelCase       =  image_std if image_std is not None else IMAGENET_DEFAULT_STD







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ = PILImageResampling.BICUBIC ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  get_size_dict(lowerCamelCase__ ,   default_to_square=lowerCamelCase__   )
      # size_dict is a dict with either keys "height" and "width" or "shortest_edge"
      if "shortest_edge" in size:
         UpperCamelCase       =  int((2_5_6 / 2_2_4) * size['''shortest_edge''']   )
         UpperCamelCase       =  get_resize_output_image_size(lowerCamelCase__ ,   size=lowerCamelCase__ ,   default_to_square=lowerCamelCase__   )
         UpperCamelCase       =  {'''height''': output_size[0], '''width''': output_size[1]}
      if "height" not in size_dict or "width" not in size_dict:
         raise ValueError(
             f'Size dict must have keys \'height\' and \'width\' or \'shortest_edge\'. Got {size_dict.keys()}'   )
      return resize(
          lowerCamelCase__ ,   size=(size_dict['''height'''], size_dict['''width''']) ,   resample=lowerCamelCase__ ,   data_format=lowerCamelCase__ ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  get_size_dict(lowerCamelCase__   )
      if "height" not in size or "width" not in size:
         raise ValueError(f'Size dict must have keys \'height\' and \'width\'. Got {size.keys()}'   )
      return center_crop(lowerCamelCase__ ,   size=(size['''height'''], size['''width''']) ,   data_format=lowerCamelCase__ ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      return rescale(lowerCamelCase__ ,   scale=lowerCamelCase__ ,   data_format=lowerCamelCase__ ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ = None ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      return normalize(lowerCamelCase__ ,   mean=lowerCamelCase__ ,   std=lowerCamelCase__ ,   data_format=lowerCamelCase__ ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = None ,   lowerCamelCase__ = ChannelDimension.FIRST ,   **lowerCamelCase__ ,   ):

      '''simple docstring'''







      UpperCamelCase       =  do_resize if do_resize is not None else self.do_resize
      UpperCamelCase       =  resample if resample is not None else self.resample
      UpperCamelCase       =  do_center_crop if do_center_crop is not None else self.do_center_crop
      UpperCamelCase       =  do_rescale if do_rescale is not None else self.do_rescale
      UpperCamelCase       =  rescale_factor if rescale_factor is not None else self.rescale_factor
      UpperCamelCase       =  do_normalize if do_normalize is not None else self.do_normalize
      UpperCamelCase       =  image_mean if image_mean is not None else self.image_mean
      UpperCamelCase       =  image_std if image_std is not None else self.image_std

      UpperCamelCase       =  size if size is not None else self.size
      UpperCamelCase       =  get_size_dict(lowerCamelCase__ ,   default_to_square=lowerCamelCase__   )
      UpperCamelCase       =  crop_size if crop_size is not None else self.crop_size
      UpperCamelCase       =  get_size_dict(lowerCamelCase__ ,   param_name='''crop_size'''   )

      UpperCamelCase       =  make_list_of_images(lowerCamelCase__   )

      if not valid_images(lowerCamelCase__   ):
         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:
         raise ValueError('''Size must be specified if do_resize is True.'''   )

      if do_center_crop and crop_size is None:
         raise ValueError('''Crop size must be specified if do_center_crop is True.'''   )

      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.
      UpperCamelCase       =  [to_numpy_array(lowerCamelCase__   ) for image in images]

      if do_resize:
         UpperCamelCase       =  [self.resize(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ) for image in images]

      if do_center_crop:
         UpperCamelCase       =  [self.center_crop(lowerCamelCase__ ,   lowerCamelCase__   ) for image in images]

      if do_rescale:
         UpperCamelCase       =  [self.rescale(lowerCamelCase__ ,   lowerCamelCase__   ) for image in images]

      if do_normalize:
         UpperCamelCase       =  [self.normalize(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ) for image in images]

      UpperCamelCase       =  [to_channel_dimension_format(lowerCamelCase__ ,   lowerCamelCase__   ) for image in images]

      UpperCamelCase       =  {'''pixel_values''': images}
      return BatchFeature(data=lowerCamelCase__ ,   tensor_type=lowerCamelCase__   )







'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list


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






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      super().__init__(**lowerCamelCase__   )

      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [vocab_size]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

      UpperCamelCase       =  self.cutoffs[0]
      UpperCamelCase       =  len(self.cutoffs   ) - 1
      UpperCamelCase       =  self.shortlist_size + self.n_clusters
      UpperCamelCase       =  keep_order

      UpperCamelCase       =  []
      UpperCamelCase       =  []







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters > 0:
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_weight'''   )
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_bias'''   )

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

            UpperCamelCase       =  self.add_weight(
                shape=(d_emb_i, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}'   )
            self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx, d_emb_i) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      super().build(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase       =  x
      if proj is not None:
         UpperCamelCase       =  tf.einsum('''ibd,ed->ibe''' ,   lowerCamelCase__ ,   lowerCamelCase__   )
      return tf.einsum('''ibd,nd->ibn''' ,   lowerCamelCase__ ,   lowerCamelCase__   ) + b







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  shape_list(lowerCamelCase__   )
      UpperCamelCase       =  tf.range(lp_size[0] ,   dtype=target.dtype   )
      UpperCamelCase       =  tf.stack([r, target] ,   1   )
      return tf.gather_nd(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=True ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      UpperCamelCase       =  0
      if self.n_clusters == 0:
         UpperCamelCase       =  self._logit(lowerCamelCase__ ,   self.out_layers[0][0] ,   self.out_layers[0][1] ,   self.out_projs[0]   )
         if target is not None:
            UpperCamelCase       =  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__ ,   logits=lowerCamelCase__   )
         UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__ ,   axis=-1   )
      else:
         UpperCamelCase       =  shape_list(lowerCamelCase__   )
         UpperCamelCase       =  []
         UpperCamelCase       =  tf.zeros(hidden_sizes[:2]   )
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            if target is not None:
               UpperCamelCase       =  (target >= l_idx) & (target < r_idx)
               UpperCamelCase       =  tf.where(lowerCamelCase__   )
               UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   ) - l_idx

            if self.div_val == 1:
               UpperCamelCase       =  self.out_layers[0][0][l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0][1][l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i][0]
               UpperCamelCase       =  self.out_layers[i][1]

            if i == 0:
               UpperCamelCase       =  tf.concat([cur_W, self.cluster_weight] ,   0   )
               UpperCamelCase       =  tf.concat([cur_b, self.cluster_bias] ,   0   )

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

      if target is not None:
         if return_mean:
            UpperCamelCase       =  tf.reduce_mean(lowerCamelCase__   )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowerCamelCase__   )

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

      return out





