


'''simple docstring'''



import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


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




def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			Union[str, Any]    ,      _lowercase      :			Optional[Any]    ,      _lowercase      :			List[str]=False						):





     '''simple docstring'''






     try:
          import torch  # noqa: F401
     except ImportError:
          logger.error(
              '''Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see'''
              ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
              ''' instructions.'''						)
          raise

     if not is_sharded:
          __UpperCAmelCase			:				Union[str, Any]													=    os.path.abspath(__snake_case						)
          logger.info(F'Loading PyTorch weights from {pt_path}'						)

          __UpperCAmelCase			:				str													=    torch.load(__snake_case    ,      map_location='''cpu'''						)
          logger.info(F'PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()						):,} parameters.'						)

          __UpperCAmelCase			:				int													=    convert_pytorch_state_dict_to_flax(__snake_case    ,      __snake_case						)
     else:
          # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
          __UpperCAmelCase			:				Dict													=    convert_pytorch_sharded_state_dict_to_flax(__snake_case    ,      __snake_case						)
     return flax_state_dict




def 			_a			(		_lowercase      :			str    ,      _lowercase      :			List[Any]    ,      _lowercase      :			List[str]    ,      _lowercase      :			Any    ,      ):





     '''simple docstring'''






     def is_key_or_prefix_key_in_dict(_lowercase      :			List[str]						) -> bool:
          return len(set(__snake_case						) & {key, (model_prefix,) + key}						) > 0

     # layer norm
     __UpperCAmelCase			:				List[Any]													=    pt_tuple_key[:-1] + ('''scale''',)
     if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__snake_case						):
          return renamed_pt_tuple_key, pt_tensor

     # batch norm layer mean
     __UpperCAmelCase			:				Optional[Any]													=    pt_tuple_key[:-1] + ('''mean''',)
     if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__snake_case						):
          return renamed_pt_tuple_key, pt_tensor

     # batch norm layer var
     __UpperCAmelCase			:				str													=    pt_tuple_key[:-1] + ('''var''',)
     if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__snake_case						):
          return renamed_pt_tuple_key, pt_tensor

     # embedding
     __UpperCAmelCase			:				Tuple													=    pt_tuple_key[:-1] + ('''embedding''',)
     if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__snake_case						):
          return renamed_pt_tuple_key, pt_tensor

     # conv layer
     __UpperCAmelCase			:				int													=    pt_tuple_key[:-1] + ('''kernel''',)
     if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__snake_case						):
          __UpperCAmelCase			:				str													=    pt_tensor.transpose(2    ,      3    ,      1    ,      0						)
          return renamed_pt_tuple_key, pt_tensor

     # linear layer
     __UpperCAmelCase			:				Optional[int]													=    pt_tuple_key[:-1] + ('''kernel''',)
     if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__snake_case						):
          __UpperCAmelCase			:				List[str]													=    pt_tensor.T
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm weight
     __UpperCAmelCase			:				Union[str, Any]													=    pt_tuple_key[:-1] + ('''weight''',)
     if pt_tuple_key[-1] == "gamma":
          return renamed_pt_tuple_key, pt_tensor

     # old PyTorch layer norm bias
     __UpperCAmelCase			:				Optional[Any]													=    pt_tuple_key[:-1] + ('''bias''',)
     if pt_tuple_key[-1] == "beta":
          return renamed_pt_tuple_key, pt_tensor

     # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
     __UpperCAmelCase			:				Tuple													=    None
     if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
          __UpperCAmelCase			:				str													=    pt_tuple_key[-2] + '''_g'''
     elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
          __UpperCAmelCase			:				Any													=    pt_tuple_key[-2] + '''_v'''
     if name is not None:
          __UpperCAmelCase			:				Optional[Any]													=    pt_tuple_key[:-3] + (name,)
          return renamed_pt_tuple_key, pt_tensor

     return pt_tuple_key, pt_tensor




def 			_a			(		_lowercase      :			Union[str, Any]    ,      _lowercase      :			Dict						):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    {k: v.numpy() for k, v in pt_state_dict.items()}

     __UpperCAmelCase			:				str													=    flax_model.base_model_prefix

     # use params dict if the model contains batch norm layers
     if "params" in flax_model.params:
          __UpperCAmelCase			:				List[str]													=    flax_model.params['''params''']
     else:
          __UpperCAmelCase			:				Dict													=    flax_model.params
     __UpperCAmelCase			:				List[Any]													=    flatten_dict(__snake_case						)

     # add batch_stats keys,values to dict
     if "batch_stats" in flax_model.params:
          __UpperCAmelCase			:				int													=    flatten_dict(flax_model.params['''batch_stats''']						)
          random_flax_state_dict.update(__snake_case						)

     __UpperCAmelCase			:				Dict													=    {}

     __UpperCAmelCase			:				List[str]													=    (model_prefix not in flax_model_params) and (
         model_prefix in {k.split('''.'''						)[0] for k in pt_state_dict.keys()}
     )
     __UpperCAmelCase			:				int													=    (model_prefix in flax_model_params) and (
         model_prefix not in {k.split('''.'''						)[0] for k in pt_state_dict.keys()}
     )

     # Need to change some parameters name to match Flax names
     for pt_key, pt_tensor in pt_state_dict.items():
          __UpperCAmelCase			:				Any													=    tuple(pt_key.split('''.'''						)						)

          # remove base model prefix if necessary
          __UpperCAmelCase			:				Tuple													=    pt_tuple_key[0] == model_prefix
          if load_model_with_head_into_base_model and has_base_model_prefix:
               __UpperCAmelCase			:				Optional[int]													=    pt_tuple_key[1:]

          # Correctly rename weight parameters
          __UpperCAmelCase							,     __UpperCAmelCase			:				Tuple													=    rename_key_and_reshape_tensor(
              __snake_case    ,      __snake_case    ,      __snake_case    ,      __snake_case						)

          # add model prefix if necessary
          __UpperCAmelCase			:				Dict													=    (model_prefix,) + flax_key in random_flax_state_dict
          if load_base_model_into_model_with_head and require_base_model_prefix:
               __UpperCAmelCase			:				Optional[Any]													=    (model_prefix,) + flax_key

          if flax_key in random_flax_state_dict:
               if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                    raise ValueError(
                        F'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                        F'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'						)

        # add batch stats if the model contains batchnorm layers
          if "batch_stats" in flax_model.params:
               if "mean" in flax_key[-1] or "var" in flax_key[-1]:
                    __UpperCAmelCase			:				Optional[int]													=    jnp.asarray(__snake_case						)
                    continue
               # remove num_batches_tracked key
               if "num_batches_tracked" in flax_key[-1]:
                    flax_state_dict.pop(__snake_case    ,      __snake_case						)
                    continue

               # also add unexpected weight so that warning is thrown
               __UpperCAmelCase			:				Optional[int]													=    jnp.asarray(__snake_case						)

          else:
               # also add unexpected weight so that warning is thrown
               __UpperCAmelCase			:				Dict													=    jnp.asarray(__snake_case						)

     return unflatten_dict(__snake_case						)




def 			_a			(		_lowercase      :			Optional[int]    ,      _lowercase      :			Optional[int]						):





     '''simple docstring'''






     import torch

     # Load the index
     __UpperCAmelCase			:				Tuple													=    {}
     for shard_file in shard_filenames:
          # load using msgpack utils
          __UpperCAmelCase			:				Tuple													=    torch.load(__snake_case						)
          __UpperCAmelCase			:				Any													=    {k: v.numpy() for k, v in pt_state_dict.items()}

          __UpperCAmelCase			:				int													=    flax_model.base_model_prefix

          # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
          if "batch_stats" in flax_model.params:
               __UpperCAmelCase			:				Union[str, Any]													=    flax_model.params['''params''']

               __UpperCAmelCase			:				Optional[Any]													=    flatten_dict(__snake_case						)
               random_flax_state_dict.update(flatten_dict(flax_model.params['''batch_stats''']						)						)
          else:
               __UpperCAmelCase			:				Union[str, Any]													=    flax_model.params
               __UpperCAmelCase			:				List[str]													=    flatten_dict(__snake_case						)

          __UpperCAmelCase			:				Any													=    (model_prefix not in flax_model_params) and (
              model_prefix in {k.split('''.'''						)[0] for k in pt_state_dict.keys()}
          )
          __UpperCAmelCase			:				Union[str, Any]													=    (model_prefix in flax_model_params) and (
              model_prefix not in {k.split('''.'''						)[0] for k in pt_state_dict.keys()}
          )
          # Need to change some parameters name to match Flax names
          for pt_key, pt_tensor in pt_state_dict.items():
               __UpperCAmelCase			:				int													=    tuple(pt_key.split('''.'''						)						)

               # remove base model prefix if necessary
               __UpperCAmelCase			:				str													=    pt_tuple_key[0] == model_prefix
               if load_model_with_head_into_base_model and has_base_model_prefix:
                    __UpperCAmelCase			:				List[str]													=    pt_tuple_key[1:]

               # Correctly rename weight parameters
               __UpperCAmelCase							,     __UpperCAmelCase			:				Union[str, Any]													=    rename_key_and_reshape_tensor(
                   __snake_case    ,      __snake_case    ,      __snake_case    ,      __snake_case						)
               # add model prefix if necessary
               __UpperCAmelCase			:				List[Any]													=    (model_prefix,) + flax_key in random_flax_state_dict
               if load_base_model_into_model_with_head and require_base_model_prefix:
                    __UpperCAmelCase			:				List[Any]													=    (model_prefix,) + flax_key

               if flax_key in random_flax_state_dict:
                    if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                         raise ValueError(
                             F'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                             F'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'						)

            # add batch stats if the model contains batchnorm layers
               if "batch_stats" in flax_model.params:
                    if "mean" in flax_key[-1]:
                         __UpperCAmelCase			:				Union[str, Any]													=    jnp.asarray(__snake_case						)
                         continue
                    if "var" in flax_key[-1]:
                         __UpperCAmelCase			:				Union[str, Any]													=    jnp.asarray(__snake_case						)
                         continue
                    # remove num_batches_tracked key
                    if "num_batches_tracked" in flax_key[-1]:
                         flax_state_dict.pop(__snake_case    ,      __snake_case						)
                         continue

                    # also add unexpected weight so that warning is thrown
                    __UpperCAmelCase			:				Optional[Any]													=    jnp.asarray(__snake_case						)

               else:
                    # also add unexpected weight so that warning is thrown
                    __UpperCAmelCase			:				List[Any]													=    jnp.asarray(__snake_case						)
     return unflatten_dict(__snake_case						)




def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			Tuple						):





     '''simple docstring'''






     __UpperCAmelCase			:				Union[str, Any]													=    os.path.abspath(__snake_case						)
     logger.info(F'Loading Flax weights from {flax_checkpoint_path}'						)

     # import correct flax class
     __UpperCAmelCase			:				Any													=    getattr(__snake_case    ,      '''Flax''' + model.__class__.__name__						)

     # load flax weight dict
     with open(__snake_case    ,      '''rb'''						) as state_f:
          try:
               __UpperCAmelCase			:				Optional[Any]													=    from_bytes(__snake_case    ,      state_f.read()						)
          except UnpicklingError:
               raise EnvironmentError(F'Unable to convert {flax_checkpoint_path} to Flax deserializable object. '						)

     return load_flax_weights_in_pytorch_model(__snake_case    ,      __snake_case						)




def 			_a			(		_lowercase      :			int    ,      _lowercase      :			Dict						):





     '''simple docstring'''






     try:
          import torch  # noqa: F401
     except ImportError:
          logger.error(
              '''Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see'''
              ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
              ''' instructions.'''						)
          raise

     # check if we have bf16 weights
     __UpperCAmelCase			:				Any													=    flatten_dict(jax.tree_util.tree_map(lambda _lowercase						: x.dtype == jnp.bfloataa    ,      __snake_case						)						).values()
     if any(__snake_case						):
          # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
          # and bf16 is not fully supported in PT yet.
          logger.warning(
              '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
              '''before loading those in PyTorch model.'''						)
          __UpperCAmelCase			:				Any													=    jax.tree_util.tree_map(
              lambda _lowercase						: params.astype(np.floataa						) if params.dtype == jnp.bfloataa else params    ,      __snake_case						)

     __UpperCAmelCase			:				Any													=    flatten_dict(__snake_case						)
     __UpperCAmelCase			:				str													=    pt_model.state_dict()

     __UpperCAmelCase			:				List[Any]													=    (pt_model.base_model_prefix in flax_state) and (
         pt_model.base_model_prefix not in {k.split('''.'''						)[0] for k in pt_model_dict.keys()}
     )
     __UpperCAmelCase			:				Any													=    (pt_model.base_model_prefix not in flax_state) and (
         pt_model.base_model_prefix in {k.split('''.'''						)[0] for k in pt_model_dict.keys()}
     )

     # keep track of unexpected & missing keys
     __UpperCAmelCase			:				Optional[int]													=    []
     __UpperCAmelCase			:				Optional[Any]													=    set(pt_model_dict.keys()						)

     for flax_key_tuple, flax_tensor in flax_state_dict.items():
          __UpperCAmelCase			:				List[Any]													=    flax_key_tuple[0] == pt_model.base_model_prefix
          __UpperCAmelCase			:				int													=    '''.'''.join((pt_model.base_model_prefix,) + flax_key_tuple						) in pt_model_dict

          # adapt flax_key to prepare for loading from/to base model only
          if load_model_with_head_into_base_model and has_base_model_prefix:
               __UpperCAmelCase			:				Optional[int]													=    flax_key_tuple[1:]
          elif load_base_model_into_model_with_head and require_base_model_prefix:
               __UpperCAmelCase			:				Tuple													=    (pt_model.base_model_prefix,) + flax_key_tuple

          # rename flax weights to PyTorch format
          if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__snake_case						) not in pt_model_dict:
               # conv layer
               __UpperCAmelCase			:				List[str]													=    flax_key_tuple[:-1] + ('''weight''',)
               __UpperCAmelCase			:				List[Any]													=    jnp.transpose(__snake_case    ,      (3, 2, 0, 1)						)
          elif flax_key_tuple[-1] == "kernel" and ".".join(__snake_case						) not in pt_model_dict:
               # linear layer
               __UpperCAmelCase			:				Dict													=    flax_key_tuple[:-1] + ('''weight''',)
               __UpperCAmelCase			:				Optional[Any]													=    flax_tensor.T
          elif flax_key_tuple[-1] in ["scale", "embedding"]:
               __UpperCAmelCase			:				str													=    flax_key_tuple[:-1] + ('''weight''',)

          # adding batch stats from flax batch norm to pt
          elif "mean" in flax_key_tuple[-1]:
               __UpperCAmelCase			:				Optional[int]													=    flax_key_tuple[:-1] + ('''running_mean''',)
          elif "var" in flax_key_tuple[-1]:
               __UpperCAmelCase			:				Dict													=    flax_key_tuple[:-1] + ('''running_var''',)

          if "batch_stats" in flax_state:
               __UpperCAmelCase			:				Optional[int]													=    '''.'''.join(flax_key_tuple[1:]						)  # Remove the params/batch_stats header
          else:
               __UpperCAmelCase			:				str													=    '''.'''.join(__snake_case						)

          # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
          __UpperCAmelCase			:				int													=    {}
          # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
          for key in pt_model_dict:
               __UpperCAmelCase			:				Any													=    key.split('''.'''						)
               __UpperCAmelCase			:				str													=    None
               if key_components[-3::2] == ["parametrizations", "original0"]:
                    __UpperCAmelCase			:				str													=    key_components[-2] + '''_g'''
               elif key_components[-3::2] == ["parametrizations", "original1"]:
                    __UpperCAmelCase			:				Union[str, Any]													=    key_components[-2] + '''_v'''
               if name is not None:
                    __UpperCAmelCase			:				int													=    key_components[:-3] + [name]
                    __UpperCAmelCase			:				str													=    '''.'''.join(__snake_case						)
                    __UpperCAmelCase			:				Tuple													=    key

          if flax_key in special_pt_names:
               __UpperCAmelCase			:				List[str]													=    special_pt_names[flax_key]

          if flax_key in pt_model_dict:
               if flax_tensor.shape != pt_model_dict[flax_key].shape:
                    raise ValueError(
                        F'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
                        F'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'						)
               else:
                    # add weight to pytorch dict
                    __UpperCAmelCase			:				Optional[int]													=    np.asarray(__snake_case						) if not isinstance(__snake_case    ,      np.ndarray						) else flax_tensor
                    __UpperCAmelCase			:				Optional[Any]													=    torch.from_numpy(__snake_case						)
                    # remove from missing keys
                    missing_keys.remove(__snake_case						)
          else:
               # weight is not expected by PyTorch model
               unexpected_keys.append(__snake_case						)

     pt_model.load_state_dict(__snake_case						)

     # re-transform missing_keys to list
     __UpperCAmelCase			:				Union[str, Any]													=    list(__snake_case						)

     if len(__snake_case						) > 0:
          logger.warning(
              '''Some weights of the Flax model were not used when initializing the PyTorch model'''
              F' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
              F' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
              ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
              F' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
              ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
              ''' FlaxBertForSequenceClassification model).'''						)
     else:
          logger.warning(F'All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'						)
     if len(__snake_case						) > 0:
          logger.warning(
              F'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
              F' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
              ''' use it for predictions and inference.'''						)
     else:
          logger.warning(
              F'All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'
              '''If your task is similar to the task the model of the checkpoint was trained on, '''
              F'you can already use {pt_model.__class__.__name__} for predictions without further training.'						)

     return pt_model








'''simple docstring'''



import numpy as np
from cva import destroyAllWindows, imread, imshow, waitKey



class 			a       :


     """simple docstring"""






     def __init__(       self  :	Union[str, Any]							,				snake_case  :	List[Any]							,				snake_case  :	int							,				snake_case  :	int    )		->							List[Any]:
          if dst_width < 0 or dst_height < 0:
               raise ValueError('''Destination width/height should be > 0'''    )

          __UpperCAmelCase			:				str													=    img
          __UpperCAmelCase			:				List[Any]													=    img.shape[1]
          __UpperCAmelCase			:				Optional[Any]													=    img.shape[0]
          __UpperCAmelCase			:				Dict													=    dst_width
          __UpperCAmelCase			:				List[str]													=    dst_height

          __UpperCAmelCase			:				Union[str, Any]													=    self.src_w / self.dst_w
          __UpperCAmelCase			:				List[str]													=    self.src_h / self.dst_h

          __UpperCAmelCase			:				Optional[int]													=    (
              np.ones((self.dst_h, self.dst_w, 3)							,				np.uinta    ) * 255
          )

     def 				lowerCamelCase__				(       self  :	Any    )		->							str:
          for i in range(self.dst_h    ):
               for j in range(self.dst_w    ):
                    __UpperCAmelCase			:				Any													=    self.img[self.get_y(snake_case    )][self.get_x(snake_case    )]

     def 				lowerCamelCase__				(       self  :	int							,				snake_case  :	int    )		->							int:
          return int(self.ratio_x * x    )






     def 				lowerCamelCase__				(       self  :	Optional[Any]							,				snake_case  :	int    )		->							int:
          return int(self.ratio_y * y    )


if __name__ == "__main__":
      __UpperCAmelCase       ,					__UpperCAmelCase    :int				       =     8_0_0, 6_0_0
      __UpperCAmelCase    :Dict				       =     imread("image_data/lena.jpg", 1)
      __UpperCAmelCase    :int				       =     NearestNeighbour(im, dst_w, dst_h)
      n.process()

      imshow(
          f"""Image resized from: {im.shape[1]}x{im.shape[0]} to {dst_w}x{dst_h}""", n.output
      )
      waitKey(0)
      destroyAllWindows()