






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 							_lowerCAmelCase:





       """simple docstring"""







       def __init__(   self    ,     _lowerCamelCase    ,     _lowerCamelCase=1_3    ,     _lowerCamelCase=3_2    ,     _lowerCamelCase=2    ,     _lowerCamelCase=3    ,     _lowerCamelCase=1_6    ,     _lowerCamelCase=[3_2, 6_4, 1_2_8]    ,     _lowerCamelCase=[1, 2, 1]    ,     _lowerCamelCase=[2, 2, 4]    ,     _lowerCamelCase=2    ,     _lowerCamelCase=2.0    ,     _lowerCamelCase=True    ,     _lowerCamelCase=0.0    ,     _lowerCamelCase=0.0    ,     _lowerCamelCase=0.1    ,     _lowerCamelCase="gelu"    ,     _lowerCamelCase=False    ,     _lowerCamelCase=True    ,     _lowerCamelCase=0.0_2    ,     _lowerCamelCase=1e-5    ,     _lowerCamelCase=True    ,     _lowerCamelCase=None    ,     _lowerCamelCase=True    ,     _lowerCamelCase=1_0    ,     _lowerCamelCase=8    ,     _lowerCamelCase=["stage1", "stage2"]    ,     _lowerCamelCase=[1, 2]    ,     ):
         UpperCamelCase_:   Tuple        =							parent
         UpperCamelCase_:   Dict        =							batch_size
         UpperCamelCase_:   List[str]        =							image_size
         UpperCamelCase_:   Tuple        =							patch_size
         UpperCamelCase_:   Tuple        =							num_channels
         UpperCamelCase_:   Dict        =							embed_dim
         UpperCamelCase_:   List[Any]        =							hidden_sizes
         UpperCamelCase_:   List[str]        =							depths
         UpperCamelCase_:   List[str]        =							num_heads
         UpperCamelCase_:   Optional[int]        =							window_size
         UpperCamelCase_:   Tuple        =							mlp_ratio
         UpperCamelCase_:   Dict        =							qkv_bias
         UpperCamelCase_:   str        =							hidden_dropout_prob
         UpperCamelCase_:   Optional[Any]        =							attention_probs_dropout_prob
         UpperCamelCase_:   int        =							drop_path_rate
         UpperCamelCase_:   Dict        =							hidden_act
         UpperCamelCase_:   List[str]        =							use_absolute_embeddings
         UpperCamelCase_:   Dict        =							patch_norm
         UpperCamelCase_:   Optional[Any]        =							layer_norm_eps
         UpperCamelCase_:   List[str]        =							initializer_range
         UpperCamelCase_:   List[Any]        =							is_training
         UpperCamelCase_:   Optional[int]        =							scope
         UpperCamelCase_:   str        =							use_labels
         UpperCamelCase_:   List[str]        =							type_sequence_label_size
         UpperCamelCase_:   Union[str, Any]        =							encoder_stride
         UpperCamelCase_:   Dict        =							out_features
         UpperCamelCase_:   str        =							out_indices





       def      _a       (   self				):
         UpperCamelCase_:   int        =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

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

         UpperCamelCase_:   Optional[Any]        =							self.get_config()

         return config, pixel_values, labels





       def      _a       (   self				):
         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      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Optional[int]        =							FocalNetModel(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   int        =							model(_lowerCamelCase				)

         UpperCamelCase_:   int        =							((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths				) - 1))
         UpperCamelCase_:   int        =							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      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   List[str]        =							FocalNetBackbone(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   Optional[int]        =							model(_lowerCamelCase				)

         # 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_:   int        =							None
         UpperCamelCase_:   List[Any]        =							FocalNetBackbone(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   Any        =							model(_lowerCamelCase				)

         # 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      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Tuple        =							FocalNetForMaskedImageModeling(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   Any        =							model(_lowerCamelCase				)
         self.parent.assertEqual(
             result.reconstruction.shape    ,     (self.batch_size, self.num_channels, self.image_size, self.image_size)				)

         # test greyscale images
         UpperCamelCase_:   List[Any]        =							1
         UpperCamelCase_:   Dict        =							FocalNetForMaskedImageModeling(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()

         UpperCamelCase_:   Optional[int]        =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
         UpperCamelCase_:   Union[str, Any]        =							model(_lowerCamelCase				)
         self.parent.assertEqual(result.reconstruction.shape    ,     (self.batch_size, 1, self.image_size, self.image_size)				)





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Tuple        =							self.type_sequence_label_size
         UpperCamelCase_:   List[Any]        =							FocalNetForImageClassification(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   str        =							model(_lowerCamelCase    ,     labels=_lowerCamelCase				)
         self.parent.assertEqual(result.logits.shape    ,     (self.batch_size, self.type_sequence_label_size)				)

         # test greyscale images
         UpperCamelCase_:   Union[str, Any]        =							1
         UpperCamelCase_:   Dict        =							FocalNetForImageClassification(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()

         UpperCamelCase_:   Dict        =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
         UpperCamelCase_:   Union[str, Any]        =							model(_lowerCamelCase				)
         self.parent.assertEqual(result.logits.shape    ,     (self.batch_size, self.type_sequence_label_size)				)





       def      _a       (   self				):
         UpperCamelCase_:   Dict        =							self.prepare_config_and_inputs()

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







@require_torch
class 							_lowerCAmelCase(			UpperCAmelCase_				,      UpperCAmelCase_				,      unittest.TestCase   ):





       """simple docstring"""







       a							:						Optional[int]										=(
           (
               FocalNetModel,
               FocalNetForImageClassification,
               FocalNetForMaskedImageModeling,
               FocalNetBackbone,
           )
           if is_torch_available()
           else ()
       )
       a							:						Any										=(
           {'''feature-extraction''': FocalNetModel, '''image-classification''': FocalNetForImageClassification}
           if is_torch_available()
           else {}
       )
       a							:						Dict										=False

       a							:						Union[str, Any]										=False
       a							:						Tuple										=False
       a							:						Optional[int]										=False
       a							:						Union[str, Any]										=False





       def      _a       (   self				):
         UpperCamelCase_:   str        =							FocalNetModelTester(self				)
         UpperCamelCase_:   Tuple        =							ConfigTester(self    ,     config_class=_lowerCamelCase    ,     embed_dim=3_7    ,     has_text_modality=_lowerCamelCase				)





       def      _a       (   self				):
         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      _a       (   self				):
         return





       def      _a       (   self				):
         UpperCamelCase_:   Optional[Any]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   List[Any]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*_lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   Optional[int]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_image_modeling(*_lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   Any        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase				)





       @unittest.skip(reason='FocalNet does not use inputs_embeds'				)
       def      _a       (   self				):
         pass





       @unittest.skip(reason='FocalNet does not use feedforward chunking'				)
       def      _a       (   self				):
         pass





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   List[Any]        =							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes[:-1]:
           UpperCamelCase_:   Union[str, Any]        =							model_class(_lowerCamelCase				)
           self.assertIsInstance(model.get_input_embeddings()    ,     (nn.Module)				)
           UpperCamelCase_:   List[str]        =							model.get_output_embeddings()
           self.assertTrue(x is None or isinstance(_lowerCamelCase    ,     nn.Linear				)				)





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Tuple        =							self.model_tester.prepare_config_and_inputs_for_common()

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

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





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Tuple        =							model_class(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()

         with torch.no_grad():
           UpperCamelCase_:   Tuple        =							model(**self._prepare_for_class(_lowerCamelCase    ,     _lowerCamelCase				)				)

         UpperCamelCase_:   Union[str, Any]        =							outputs.hidden_states

         UpperCamelCase_:   Tuple        =							getattr(
             self.model_tester    ,     'expected_num_hidden_layers'    ,     len(self.model_tester.depths				) + 1				)
         self.assertEqual(len(_lowerCamelCase				)    ,     _lowerCamelCase				)

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

         UpperCamelCase_:   int        =							(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_:   Dict        =							outputs.reshaped_hidden_states
         self.assertEqual(len(_lowerCamelCase				)    ,     _lowerCamelCase				)

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





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Union[str, Any]        =							self.model_tester.prepare_config_and_inputs_for_common()

         UpperCamelCase_:   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_:   int        =							True
           self.check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           UpperCamelCase_:   Optional[Any]        =							True

           self.check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Dict        =							self.model_tester.prepare_config_and_inputs_for_common()
         UpperCamelCase_:   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_:   int        =							(
             config.patch_size
             if isinstance(config.patch_size    ,     collections.abc.Iterable				)
             else (config.patch_size, config.patch_size)
         )

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

         for model_class in self.all_model_classes[:-1]:
           UpperCamelCase_:   Dict        =							True
           self.check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     (padded_height, padded_width)				)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           UpperCamelCase_:   Optional[Any]        =							True
           self.check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     (padded_height, padded_width)				)





       @slow
       def      _a       (   self				):
         for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase_:   List[Any]        =							FocalNetModel.from_pretrained(_lowerCamelCase				)
           self.assertIsNotNone(_lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   Dict        =							self.model_tester.prepare_config_and_inputs_for_common()

         UpperCamelCase_:   Dict        =							_config_zero_init(_lowerCamelCase				)
         for model_class in self.all_model_classes:
           UpperCamelCase_:   List[str]        =							model_class(config=_lowerCamelCase				)
           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 							_lowerCAmelCase(			unittest.TestCase   ):





       """simple docstring"""







       @cached_property
       def      _a       (   self				):
         # TODO update organization
         return AutoImageProcessor.from_pretrained('microsoft/focalnet-tiny'				) if is_vision_available() else None





       @slow
       def      _a       (   self				):
         UpperCamelCase_:   Optional[int]        =							FocalNetForImageClassification.from_pretrained('microsoft/focalnet-tiny'				).to(_lowerCamelCase				)
         UpperCamelCase_:   Optional[Any]        =							self.default_image_processor

         UpperCamelCase_:   str        =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
         UpperCamelCase_:   str        =							image_processor(images=_lowerCamelCase    ,     return_tensors='pt'				).to(_lowerCamelCase				)

         # forward pass
         with torch.no_grad():
           UpperCamelCase_:   List[str]        =							model(**_lowerCamelCase				)

         # verify the logits
         UpperCamelCase_:   Optional[int]        =							torch.Size((1, 1_0_0_0)				)
         self.assertEqual(outputs.logits.shape    ,     _lowerCamelCase				)
         UpperCamelCase_:   Optional[int]        =							torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]				).to(_lowerCamelCase				)
         self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,     _lowerCamelCase    ,     atol=1e-4				)				)
         self.assertTrue(outputs.logits.argmax(dim=-1				).item()    ,     2_8_1				)







@require_torch
class 							_lowerCAmelCase(			UpperCAmelCase_				,      unittest.TestCase   ):





       """simple docstring"""







       a							:						Optional[Any]										=(FocalNetBackbone,) if is_torch_available() else ()
       a							:						List[str]										=FocalNetConfig

       a							:						List[str]										=False





       def      _a       (   self				):
         UpperCamelCase_:   Any        =							FocalNetModelTester(self				)
'''simple docstring'''




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__UpperCAmelCase					=object()

# For specifying empty leaf dict `{}`
__UpperCAmelCase					=object()






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       List[Any]:
    __lowerCamelCase            =    tuple((re.compile(x + '''$'''							) for x in qs)							)
    for i in range(len(UpperCamelCase__							) - len(UpperCamelCase__							) + 1							):
        __lowerCamelCase            =    [x.match(UpperCamelCase__							) for x, y in zip(UpperCamelCase__							,    ks[i:]							)]
        if matches and all(UpperCamelCase__							):
            return True
    return False






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Optional[int]:
    def replace(UpperCamelCase__							,    UpperCamelCase__							):
        for rule, replacement in rules:
            if _match(UpperCamelCase__							,    UpperCamelCase__							):
                return replacement
        return val

    return replace






def __lowerCAmelCase					(				)      ->       Dict:
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        (("transformer", "wte", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("attention", "out_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("mlp", "c_fc", "bias"), P('''mp'''							)),
        (("mlp", "c_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       int:
    __lowerCamelCase            =    _get_partition_rules()
    __lowerCamelCase            =    _replacement_rules(UpperCamelCase__							)
    __lowerCamelCase            =    {k: _unmatched for k in flatten_dict(UpperCamelCase__							)}
    __lowerCamelCase            =    {k: replace(UpperCamelCase__							,    UpperCamelCase__							) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(UpperCamelCase__							)							)
