
"""simple docstring"""





import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        FocalNetBackbone,
        FocalNetForImageClassification,
        FocalNetForMaskedImageModeling,
        FocalNetModel,
    )
    from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor


class A_	:




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=13				,   lowercase_=32				,   lowercase_=2				,   lowercase_=3				,   lowercase_=16				,   lowercase_=[32, 64, 128]				,   lowercase_=[1, 2, 1]				,   lowercase_=[2, 2, 4]				,   lowercase_=2				,   lowercase_=2.0				,   lowercase_=True				,   lowercase_=0.0				,   lowercase_=0.0				,   lowercase_=0.1				,   lowercase_="gelu"				,   lowercase_=False				,   lowercase_=True				,   lowercase_=0.02				,   lowercase_=1E-5				,   lowercase_=True				,   lowercase_=None				,   lowercase_=True				,   lowercase_=10				,   lowercase_=8				,   lowercase_=["stage1", "stage2"]				,   lowercase_=[1, 2]				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    parent
     UpperCAmelCase_					:					Optional[int]			     =    batch_size
     UpperCAmelCase_					:					Optional[int]			     =    image_size
     UpperCAmelCase_					:					Union[str, Any]			     =    patch_size
     UpperCAmelCase_					:					Union[str, Any]			     =    num_channels
     UpperCAmelCase_					:					Optional[int]			     =    embed_dim
     UpperCAmelCase_					:					Dict			     =    hidden_sizes
     UpperCAmelCase_					:					str			     =    depths
     UpperCAmelCase_					:					Any			     =    num_heads
     UpperCAmelCase_					:					Optional[Any]			     =    window_size
     UpperCAmelCase_					:					List[str]			     =    mlp_ratio
     UpperCAmelCase_					:					Any			     =    qkv_bias
     UpperCAmelCase_					:					Dict			     =    hidden_dropout_prob
     UpperCAmelCase_					:					List[str]			     =    attention_probs_dropout_prob
     UpperCAmelCase_					:					str			     =    drop_path_rate
     UpperCAmelCase_					:					Union[str, Any]			     =    hidden_act
     UpperCAmelCase_					:					Optional[Any]			     =    use_absolute_embeddings
     UpperCAmelCase_					:					Optional[Any]			     =    patch_norm
     UpperCAmelCase_					:					Any			     =    layer_norm_eps
     UpperCAmelCase_					:					Optional[int]			     =    initializer_range
     UpperCAmelCase_					:					str			     =    is_training
     UpperCAmelCase_					:					Any			     =    scope
     UpperCAmelCase_					:					str			     =    use_labels
     UpperCAmelCase_					:					Tuple			     =    type_sequence_label_size
     UpperCAmelCase_					:					str			     =    encoder_stride
     UpperCAmelCase_					:					Any			     =    out_features
     UpperCAmelCase_					:					List[str]			     =    out_indices







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

     UpperCAmelCase_					:					Any			     =    None
     if self.use_labels:
        UpperCAmelCase_					:					Tuple			     =    ids_tensor([self.batch_size]				,   self.type_sequence_label_size  )

     UpperCAmelCase_					:					str			     =    self.get_config()

     return config, pixel_values, labels







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return FocalNetConfig(
         image_size=self.image_size				,   patch_size=self.patch_size				,   num_channels=self.num_channels				,   embed_dim=self.embed_dim				,   hidden_sizes=self.hidden_sizes				,   depths=self.depths				,   num_heads=self.num_heads				,   window_size=self.window_size				,   mlp_ratio=self.mlp_ratio				,   qkv_bias=self.qkv_bias				,   hidden_dropout_prob=self.hidden_dropout_prob				,   attention_probs_dropout_prob=self.attention_probs_dropout_prob				,   drop_path_rate=self.drop_path_rate				,   hidden_act=self.hidden_act				,   use_absolute_embeddings=self.use_absolute_embeddings				,   path_norm=self.patch_norm				,   layer_norm_eps=self.layer_norm_eps				,   initializer_range=self.initializer_range				,   encoder_stride=self.encoder_stride				,   out_features=self.out_features				,   out_indices=self.out_indices				,   )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    FocalNetModel(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					Optional[Any]			     =    model(lowercase_  )

     UpperCAmelCase_					:					Tuple			     =    ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths  ) - 1))
     UpperCAmelCase_					:					Any			     =    int(config.embed_dim * 2 ** (len(config.depths  ) - 1)  )

     self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, expected_seq_len, expected_dim)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    FocalNetBackbone(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					int			     =    model(lowercase_  )

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps  )				,   len(config.out_features  )  )
     self.parent.assertListEqual(list(result.feature_maps[0].shape  )				,   [self.batch_size, self.image_size, 8, 8]  )

     # verify channels
     self.parent.assertEqual(len(model.channels  )				,   len(config.out_features  )  )
     self.parent.assertListEqual(model.channels				,   config.hidden_sizes[:-1]  )

     # verify backbone works with out_features=None
     UpperCAmelCase_					:					List[Any]			     =    None
     UpperCAmelCase_					:					str			     =    FocalNetBackbone(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					str			     =    model(lowercase_  )

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps  )				,   1  )
     self.parent.assertListEqual(list(result.feature_maps[0].shape  )				,   [self.batch_size, self.image_size * 2, 4, 4]  )

     # verify channels
     self.parent.assertEqual(len(model.channels  )				,   1  )
     self.parent.assertListEqual(model.channels				,   [config.hidden_sizes[-1]]  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    FocalNetForMaskedImageModeling(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					Tuple			     =    model(lowercase_  )
     self.parent.assertEqual(
         result.reconstruction.shape				,   (self.batch_size, self.num_channels, self.image_size, self.image_size)  )

     # test greyscale images
     UpperCAmelCase_					:					List[str]			     =    1
     UpperCAmelCase_					:					Any			     =    FocalNetForMaskedImageModeling(lowercase_  )
     model.to(lowercase_  )
     model.eval()

     UpperCAmelCase_					:					Optional[int]			     =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
     UpperCAmelCase_					:					Dict			     =    model(lowercase_  )
     self.parent.assertEqual(result.reconstruction.shape				,   (self.batch_size, 1, self.image_size, self.image_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    self.type_sequence_label_size
     UpperCAmelCase_					:					Optional[int]			     =    FocalNetForImageClassification(lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					Tuple			     =    model(lowercase_				,   labels=lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)  )

     # test greyscale images
     UpperCAmelCase_					:					Optional[int]			     =    1
     UpperCAmelCase_					:					Optional[int]			     =    FocalNetForImageClassification(lowercase_  )
     model.to(lowercase_  )
     model.eval()

     UpperCAmelCase_					:					int			     =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
     UpperCAmelCase_					:					str			     =    model(lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    self.prepare_config_and_inputs()

     UpperCAmelCase_					:					List[str]			     =    config_and_inputs
     UpperCAmelCase_					:					Optional[int]			     =    {'''pixel_values''': pixel_values}
     return config, inputs_dict




@require_torch
class A_	(UpperCAmelCase__				,UpperCAmelCase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Optional[Any]       =		(
      (
          FocalNetModel,
          FocalNetForImageClassification,
          FocalNetForMaskedImageModeling,
          FocalNetBackbone,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__			:						List[str]       =		(
      {"""feature-extraction""": FocalNetModel, """image-classification""": FocalNetForImageClassification}
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__			:						Dict       =		False

  SCREAMING_SNAKE_CASE__			:						Dict       =		False
  SCREAMING_SNAKE_CASE__			:						Tuple       =		False
  SCREAMING_SNAKE_CASE__			:						Any       =		False
  SCREAMING_SNAKE_CASE__			:						str       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    FocalNetModelTester(self  )
     UpperCAmelCase_					:					List[str]			     =    ConfigTester(self				,   config_class=lowercase_				,   embed_dim=37				,   has_text_modality=lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					str			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_  )







  @unittest.skip(reason="FocalNet does not use inputs_embeds"  )
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     pass







  @unittest.skip(reason="FocalNet does not use feedforward chunking"  )
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     pass







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes[:-1]:
        UpperCAmelCase_					:					Any			     =    model_class(lowercase_  )
        self.assertIsInstance(model.get_input_embeddings()				,   (nn.Module)  )
        UpperCAmelCase_					:					Optional[int]			     =    model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowercase_				,   nn.Linear  )  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes[:-1]:
        UpperCAmelCase_					:					List[Any]			     =    model_class(lowercase_  )
        UpperCAmelCase_					:					List[Any]			     =    inspect.signature(model.forward  )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCAmelCase_					:					Optional[int]			     =    [*signature.parameters.keys()]

        UpperCAmelCase_					:					List[Any]			     =    ['''pixel_values''']
        self.assertListEqual(arg_names[:1]				,   lowercase_  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    model_class(lowercase_  )
     model.to(lowercase_  )
     model.eval()

     with torch.no_grad():
        UpperCAmelCase_					:					Optional[Any]			     =    model(**self._prepare_for_class(lowercase_				,   lowercase_  )  )

     UpperCAmelCase_					:					List[str]			     =    outputs.hidden_states

     UpperCAmelCase_					:					List[Any]			     =    getattr(
         self.model_tester				,   "expected_num_hidden_layers"				,   len(self.model_tester.depths  ) + 1  )
     self.assertEqual(len(lowercase_  )				,   lowercase_  )

     # FocalNet has a different seq_length
     UpperCAmelCase_					:					List[Any]			     =    (
         config.patch_size
         if isinstance(config.patch_size				,   collections.abc.Iterable  )
         else (config.patch_size, config.patch_size)
     )

     UpperCAmelCase_					:					List[Any]			     =    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

     self.assertListEqual(
         list(hidden_states[0].shape[-2:]  )				,   [num_patches, self.model_tester.embed_dim]				,   )

     UpperCAmelCase_					:					Optional[int]			     =    outputs.reshaped_hidden_states
     self.assertEqual(len(lowercase_  )				,   lowercase_  )

     UpperCAmelCase_					:					List[Any]			     =    reshaped_hidden_states[0].shape
     UpperCAmelCase_					:					int			     =    (
         reshaped_hidden_states[0].view(lowercase_				,   lowercase_				,   height * width  ).permute(0				,   2				,   1  )
     )
     self.assertListEqual(
         list(reshaped_hidden_states.shape[-2:]  )				,   [num_patches, self.model_tester.embed_dim]				,   )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    self.model_tester.prepare_config_and_inputs_for_common()

     UpperCAmelCase_					:					Optional[int]			     =    (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size				,   collections.abc.Iterable  )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes[:-1]:
        UpperCAmelCase_					:					List[Any]			     =    True
        self.check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase_					:					List[str]			     =    True

        self.check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_				,   lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[Any]			     =    self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_					:					List[str]			     =    3

     UpperCAmelCase_					:					Any			     =    (
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size				,   collections.abc.Iterable  )
         else (self.model_tester.image_size, self.model_tester.image_size)
     )
     UpperCAmelCase_					:					Tuple			     =    (
         config.patch_size
         if isinstance(config.patch_size				,   collections.abc.Iterable  )
         else (config.patch_size, config.patch_size)
     )

     UpperCAmelCase_					:					Any			     =    image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
     UpperCAmelCase_					:					List[str]			     =    image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

     for model_class in self.all_model_classes[:-1]:
        UpperCAmelCase_					:					Any			     =    True
        self.check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_				,   (padded_height, padded_width)  )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase_					:					Dict			     =    True
        self.check_hidden_states_output(lowercase_				,   lowercase_				,   lowercase_				,   (padded_height, padded_width)  )







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""


     for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase_					:					Union[str, Any]			     =    FocalNetModel.from_pretrained(lowercase_  )
        self.assertIsNotNone(lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    self.model_tester.prepare_config_and_inputs_for_common()

     UpperCAmelCase_					:					Dict			     =    _config_zero_init(lowercase_  )
     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					List[str]			     =    model_class(config=lowercase_  )
        for name, param in model.named_parameters():
           if "embeddings" not in name and param.requires_grad:
              self.assertIn(
                  ((param.data.mean() * 1E9).round() / 1E9).item()				,   [0.0, 1.0]				,   msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""				,   )




@require_vision
@require_torch
class A_	(unittest.TestCase      ):




  '''simple docstring'''







  @cached_property
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     return AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny"  ) if is_vision_available() else None







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    FocalNetForImageClassification.from_pretrained("microsoft/focalnet-tiny"  ).to(lowercase_  )
     UpperCAmelCase_					:					List[str]			     =    self.default_image_processor

     UpperCAmelCase_					:					List[str]			     =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
     UpperCAmelCase_					:					str			     =    image_processor(images=lowercase_				,   return_tensors="pt"  ).to(lowercase_  )

     # forward pass
     with torch.no_grad():
        UpperCAmelCase_					:					Union[str, Any]			     =    model(**lowercase_  )

     # verify the logits
     UpperCAmelCase_					:					Optional[Any]			     =    torch.Size((1, 1000)  )
     self.assertEqual(outputs.logits.shape				,   lowercase_  )
     UpperCAmelCase_					:					List[str]			     =    torch.tensor([0.21_66, -0.43_68, 0.21_91]  ).to(lowercase_  )
     self.assertTrue(torch.allclose(outputs.logits[0, :3]				,   lowercase_				,   atol=1E-4  )  )
     self.assertTrue(outputs.logits.argmax(dim=-1  ).item()				,   281  )




@require_torch
class A_	(UpperCAmelCase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						List[Any]       =		(FocalNetBackbone,) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE__			:						Optional[Any]       =		FocalNetConfig

  SCREAMING_SNAKE_CASE__			:						Union[str, Any]       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[Any]			     =    FocalNetModelTester(self  )










"""simple docstring"""






from pickle import UnpicklingError

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

from ..utils import logging


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



def 					_A			(__a    ,				__a							) ->       Tuple:
    """simple docstring"""






    try:
        with open(__a    ,				'''rb'''							) as flax_state_f:
            SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		from_bytes(__a    ,				flax_state_f.read()							)
    except UnpicklingError as e:
        try:
            with open(__a							) as f:
                if f.read().startswith('''version'''							):
                    raise OSError(
                        '''You seem to have cloned a repository without having git-lfs installed. Please'''
                        ''' install git-lfs and run `git lfs install` followed by `git lfs pull` in the'''
                        ''' folder you cloned.'''							)
                else:
                    raise ValueError from e
        except (UnicodeDecodeError, ValueError):
            raise EnvironmentError(f'Unable to convert {model_file} to Flax deserializable object. '							)

    return load_flax_weights_in_pytorch_model(__a    ,				__a							)





def 					_A			(__a    ,				__a							) ->       Tuple:
    """simple docstring"""






    try:
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            '''Loading 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
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		flatten_dict(jax.tree_util.tree_map(lambda __a							: x.dtype == jnp.bfloataa    ,				__a							)							).values()
    if any(__a							):
        # convert all weights to fp32 if they 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.'''							)
        SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		jax.tree_util.tree_map(
            lambda __a							: params.astype(np.floataa							) if params.dtype == jnp.bfloataa else params    ,				__a							)

    SCREAMING_SNAKE_CASE_      :  int    						=		''''''

    SCREAMING_SNAKE_CASE_      :  str    						=		flatten_dict(__a    ,				sep='''.'''							)
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		pt_model.state_dict()

    # keep track of unexpected & missing keys
    SCREAMING_SNAKE_CASE_      :  str    						=		[]
    SCREAMING_SNAKE_CASE_      :  Any    						=		set(pt_model_dict.keys()							)

    for flax_key_tuple, flax_tensor in flax_state_dict.items():
        SCREAMING_SNAKE_CASE_      :  Any    						=		flax_key_tuple.split('''.'''							)

        if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
            SCREAMING_SNAKE_CASE_      :  Any    						=		flax_key_tuple_array[:-1] + ['''weight''']
            SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		jnp.transpose(__a    ,				(3, 2, 0, 1)							)
        elif flax_key_tuple_array[-1] == "kernel":
            SCREAMING_SNAKE_CASE_      :  Tuple    						=		flax_key_tuple_array[:-1] + ['''weight''']
            SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		flax_tensor.T
        elif flax_key_tuple_array[-1] == "scale":
            SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		flax_key_tuple_array[:-1] + ['''weight''']

        if "time_embedding" not in flax_key_tuple_array:
            for i, flax_key_tuple_string in enumerate(__a							):
                SCREAMING_SNAKE_CASE_      :  List[str]    						=		(
                    flax_key_tuple_string.replace('''_0'''    ,				'''.0'''							)
                    .replace('''_1'''    ,				'''.1'''							)
                    .replace('''_2'''    ,				'''.2'''							)
                    .replace('''_3'''    ,				'''.3'''							)
                    .replace('''_4'''    ,				'''.4'''							)
                    .replace('''_5'''    ,				'''.5'''							)
                    .replace('''_6'''    ,				'''.6'''							)
                    .replace('''_7'''    ,				'''.7'''							)
                    .replace('''_8'''    ,				'''.8'''							)
                    .replace('''_9'''    ,				'''.9'''							)
                )

        SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		'''.'''.join(__a							)

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

    pt_model.load_state_dict(__a							)

    # re-transform missing_keys to list
    SCREAMING_SNAKE_CASE_      :  int    						=		list(__a							)

    if len(__a							) > 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).'''							)
    if len(__a							) > 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.'''							)

    return pt_model


