
"""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


_a		       =							logging.get_logger(__name__)







def 	__a     (		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase=False  ):
   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_					:					Any			     =    os.path.abspath(__lowerCamelCase  )
      logger.info(f"""Loading PyTorch weights from {pt_path}"""  )

      UpperCAmelCase_					:					Any			     =    torch.load(__lowerCamelCase,		map_location="cpu"  )
      logger.info(f"""PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()  ):,} parameters."""  )

      UpperCAmelCase_					:					Union[str, Any]			     =    convert_pytorch_state_dict_to_flax(__lowerCamelCase,		__lowerCamelCase  )
   else:
      # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
      UpperCAmelCase_					:					Any			     =    convert_pytorch_sharded_state_dict_to_flax(__lowerCamelCase,		__lowerCamelCase  )
   return flax_state_dict







def 	__a     (		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		):

   def is_key_or_prefix_key_in_dict(__lowerCamelCase  ) -> bool:
      return len(set(__lowerCamelCase  ) & {key, (model_prefix,) + key}  ) > 0

   # layer norm
   UpperCAmelCase_					:					Optional[Any]			     =    pt_tuple_key[:-1] + ("scale",)
   if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCamelCase  ):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer mean
   UpperCAmelCase_					:					List[Any]			     =    pt_tuple_key[:-1] + ("mean",)
   if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCamelCase  ):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer var
   UpperCAmelCase_					:					List[Any]			     =    pt_tuple_key[:-1] + ("var",)
   if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCamelCase  ):
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   UpperCAmelCase_					:					Dict			     =    pt_tuple_key[:-1] + ("embedding",)
   if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCamelCase  ):
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   UpperCAmelCase_					:					List[Any]			     =    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(__lowerCamelCase  ):
      UpperCAmelCase_					:					Optional[int]			     =    pt_tensor.transpose(2,		3,		1,		0  )
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   UpperCAmelCase_					:					str			     =    pt_tuple_key[:-1] + ("kernel",)
   if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCamelCase  ):
      UpperCAmelCase_					:					List[Any]			     =    pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   UpperCAmelCase_					:					Tuple			     =    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[int]			     =    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_					:					Tuple			     =    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_					:					Any			     =    pt_tuple_key[:-3] + (name,)
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor







def 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):
   # convert pytorch tensor to numpy
   UpperCAmelCase_					:					List[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_					:					int			     =    flax_model.params["params"]
   else:
      UpperCAmelCase_					:					int			     =    flax_model.params
   UpperCAmelCase_					:					Tuple			     =    flatten_dict(__lowerCamelCase  )

   # add batch_stats keys,values to dict
   if "batch_stats" in flax_model.params:
      UpperCAmelCase_					:					Union[str, Any]			     =    flatten_dict(flax_model.params["batch_stats"]  )
      random_flax_state_dict.update(__lowerCamelCase  )

   UpperCAmelCase_					:					int			     =    {}

   UpperCAmelCase_					:					Optional[int]			     =    (model_prefix not in flax_model_params) and (
       model_prefix in {k.split("."  )[0] for k in pt_state_dict.keys()}
   )
   UpperCAmelCase_					:					Optional[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_					:					Any			     =    tuple(pt_key.split("."  )  )

      # remove base model prefix if necessary
      UpperCAmelCase_					:					Any			     =    pt_tuple_key[0] == model_prefix
      if load_model_with_head_into_base_model and has_base_model_prefix:
         UpperCAmelCase_					:					Any			     =    pt_tuple_key[1:]

      # Correctly rename weight parameters
      UpperCAmelCase_     ,   UpperCAmelCase_					:					List[Any]			     =    rename_key_and_reshape_tensor(
          __lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase  )

      # add model prefix if necessary
      UpperCAmelCase_					:					List[str]			     =    (model_prefix,) + flax_key in random_flax_state_dict
      if load_base_model_into_model_with_head and require_base_model_prefix:
         UpperCAmelCase_					:					Tuple			     =    (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_					:					Union[str, Any]			     =    jnp.asarray(__lowerCamelCase  )
            continue
         # remove num_batches_tracked key
         if "num_batches_tracked" in flax_key[-1]:
            flax_state_dict.pop(__lowerCamelCase,		__lowerCamelCase  )
            continue

         # also add unexpected weight so that warning is thrown
         UpperCAmelCase_					:					Tuple			     =    jnp.asarray(__lowerCamelCase  )

      else:
         # also add unexpected weight so that warning is thrown
         UpperCAmelCase_					:					Any			     =    jnp.asarray(__lowerCamelCase  )

   return unflatten_dict(__lowerCamelCase  )







def 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):
   import torch

   # Load the index
   UpperCAmelCase_					:					Any			     =    {}
   for shard_file in shard_filenames:
      # load using msgpack utils
      UpperCAmelCase_					:					List[str]			     =    torch.load(__lowerCamelCase  )
      UpperCAmelCase_					:					List[str]			     =    {k: v.numpy() for k, v in pt_state_dict.items()}

      UpperCAmelCase_					:					Any			     =    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_					:					Dict			     =    flax_model.params["params"]

         UpperCAmelCase_					:					Optional[Any]			     =    flatten_dict(__lowerCamelCase  )
         random_flax_state_dict.update(flatten_dict(flax_model.params["batch_stats"]  )  )
      else:
         UpperCAmelCase_					:					Union[str, Any]			     =    flax_model.params
         UpperCAmelCase_					:					str			     =    flatten_dict(__lowerCamelCase  )

      UpperCAmelCase_					:					List[Any]			     =    (model_prefix not in flax_model_params) and (
          model_prefix in {k.split("."  )[0] for k in pt_state_dict.keys()}
      )
      UpperCAmelCase_					:					List[str]			     =    (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_					:					int			     =    pt_tuple_key[0] == model_prefix
         if load_model_with_head_into_base_model and has_base_model_prefix:
            UpperCAmelCase_					:					Union[str, Any]			     =    pt_tuple_key[1:]

         # Correctly rename weight parameters
         UpperCAmelCase_     ,   UpperCAmelCase_					:					str			     =    rename_key_and_reshape_tensor(
             __lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase,		__lowerCamelCase  )
         # 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]:
               UpperCAmelCase_					:					Tuple			     =    jnp.asarray(__lowerCamelCase  )
               continue
            if "var" in flax_key[-1]:
               UpperCAmelCase_					:					str			     =    jnp.asarray(__lowerCamelCase  )
               continue
            # remove num_batches_tracked key
            if "num_batches_tracked" in flax_key[-1]:
               flax_state_dict.pop(__lowerCamelCase,		__lowerCamelCase  )
               continue

            # also add unexpected weight so that warning is thrown
            UpperCAmelCase_					:					Tuple			     =    jnp.asarray(__lowerCamelCase  )

         else:
            # also add unexpected weight so that warning is thrown
            UpperCAmelCase_					:					str			     =    jnp.asarray(__lowerCamelCase  )
   return unflatten_dict(__lowerCamelCase  )







def 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):
   UpperCAmelCase_					:					Dict			     =    os.path.abspath(__lowerCamelCase  )
   logger.info(f"""Loading Flax weights from {flax_checkpoint_path}"""  )

   # import correct flax class
   UpperCAmelCase_					:					List[Any]			     =    getattr(__lowerCamelCase,		"Flax" + model.__class__.__name__  )

   # load flax weight dict
   with open(__lowerCamelCase,		"rb"  ) as state_f:
      try:
         UpperCAmelCase_					:					List[Any]			     =    from_bytes(__lowerCamelCase,		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(__lowerCamelCase,		__lowerCamelCase  )







def 	__a     (		__lowerCamelCase,		__lowerCamelCase  ):

   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_					:					Optional[Any]			     =    flatten_dict(jax.tree_util.tree_map(lambda __lowerCamelCase  : x.dtype == jnp.bfloataa,		__lowerCamelCase  )  ).values()
   if any(__lowerCamelCase  ):
      # 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_					:					str			     =    jax.tree_util.tree_map(
          lambda __lowerCamelCase  : params.astype(np.floataa  ) if params.dtype == jnp.bfloataa else params,		__lowerCamelCase  )

   UpperCAmelCase_					:					Dict			     =    flatten_dict(__lowerCamelCase  )
   UpperCAmelCase_					:					Dict			     =    pt_model.state_dict()

   UpperCAmelCase_					:					Tuple			     =    (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_					:					Optional[int]			     =    (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_					:					List[Any]			     =    []
   UpperCAmelCase_					:					List[Any]			     =    set(pt_model_dict.keys()  )

   for flax_key_tuple, flax_tensor in flax_state_dict.items():
      UpperCAmelCase_					:					List[str]			     =    flax_key_tuple[0] == pt_model.base_model_prefix
      UpperCAmelCase_					:					Any			     =    ".".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_					:					Dict			     =    flax_key_tuple[1:]
      elif load_base_model_into_model_with_head and require_base_model_prefix:
         UpperCAmelCase_					:					Optional[int]			     =    (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(__lowerCamelCase  ) not in pt_model_dict:
         # conv layer
         UpperCAmelCase_					:					List[Any]			     =    flax_key_tuple[:-1] + ("weight",)
         UpperCAmelCase_					:					Optional[Any]			     =    jnp.transpose(__lowerCamelCase,		(3, 2, 0, 1)  )
      elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCamelCase  ) not in pt_model_dict:
         # linear layer
         UpperCAmelCase_					:					List[Any]			     =    flax_key_tuple[:-1] + ("weight",)
         UpperCAmelCase_					:					str			     =    flax_tensor.T
      elif flax_key_tuple[-1] in ["scale", "embedding"]:
         UpperCAmelCase_					:					Optional[Any]			     =    flax_key_tuple[:-1] + ("weight",)

      # adding batch stats from flax batch norm to pt
      elif "mean" in flax_key_tuple[-1]:
         UpperCAmelCase_					:					Tuple			     =    flax_key_tuple[:-1] + ("running_mean",)
      elif "var" in flax_key_tuple[-1]:
         UpperCAmelCase_					:					str			     =    flax_key_tuple[:-1] + ("running_var",)

      if "batch_stats" in flax_state:
         UpperCAmelCase_					:					Any			     =    ".".join(flax_key_tuple[1:]  )  # Remove the params/batch_stats header
      else:
         UpperCAmelCase_					:					List[Any]			     =    ".".join(__lowerCamelCase  )

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

      if flax_key in special_pt_names:
         UpperCAmelCase_					:					Union[str, Any]			     =    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_					:					Tuple			     =    np.asarray(__lowerCamelCase  ) if not isinstance(__lowerCamelCase,		np.ndarray  ) else flax_tensor
            UpperCAmelCase_					:					Optional[int]			     =    torch.from_numpy(__lowerCamelCase  )
            # remove from missing keys
            missing_keys.remove(__lowerCamelCase  )
      else:
         # weight is not expected by PyTorch model
         unexpected_keys.append(__lowerCamelCase  )

   pt_model.load_state_dict(__lowerCamelCase  )

   # re-transform missing_keys to list
   UpperCAmelCase_					:					Any			     =    list(__lowerCamelCase  )

   if len(__lowerCamelCase  ) > 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(__lowerCamelCase  ) > 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"""




from __future__ import annotations

from functools import lru_cache
from math import ceil

_A			:       Optional[Any]   		=       1_00

_A			:       Optional[int]   		=       set(range(3, NUM_PRIMES, 2))
primes.add(2)
_A			:       int

for prime in range(3, ceil(NUM_PRIMES**0.5), 2):
   if prime not in primes:
      continue
   primes.difference_update(set(range(prime * prime, NUM_PRIMES, prime)))


@lru_cache(maxsize=100  )
def   __magic_name__	(      __snake_case				:       int  ) ->				set[int]:
       if number_to_partition < 0:
              return set()
       elif number_to_partition == 0:
              return {1}

       lowercase    :				set[int]       		=							set()
       lowercase    :				int
       lowercase    :				int

       for prime in primes:
              if prime > number_to_partition:
                     continue
              for sub in partition(number_to_partition - prime  ):
                     ret.add(sub * prime  )

       return ret


def   __magic_name__	(      __snake_case				:       int = 5000  ) ->				int | None:
       for number_to_partition in range(1     , __snake_case  ):
              if len(partition(__snake_case  )  ) > number_unique_partitions:
                     return number_to_partition
       return None


if __name__ == "__main__":
   print(F"{solution() = }")


