

"""simple docstring"""

import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
    import os

    import jax.numpy as jnp
    from jax import jit

    from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
    from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

    _UpperCamelCase		     =    """0.12"""  # assumed parallelism: 8


if is_torch_available():
    import torch



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__					,					lowercase__=None   )			->						List[Any]:
 if rng is None:
  lowerCAmelCase__   :				int 			=				random.Random()

 lowerCAmelCase__   :				int 			=				1
 for dim in shape:
  total_dims *= dim

 lowerCAmelCase__   :				List[str] 			=				[]
 for _ in range(lowercase__   ):
  values.append(rng.randint(0					,					vocab_size - 1   )   )

 lowerCAmelCase__   :				Optional[Any] 			=				np.array(lowercase__					,					dtype=jnp.intaa   ).reshape(lowercase__   )

 return output



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__=None   )			->						Optional[int]:
 lowerCAmelCase__   :				int 			=				ids_tensor(lowercase__					,					vocab_size=2					,					rng=lowercase__   )
 # make sure that at least one token is attended to for each batch
 lowerCAmelCase__   :				List[Any] 			=				1
 return attn_mask






@require_flax
class 			__a					:





       """simple docstring"""


       __UpperCamelCase					:      Dict    =						None
       __UpperCamelCase					:      Dict    =						()






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__   :				int 			=				self.model_tester.prepare_config_and_inputs_for_common()

        # cut to half length & take max batch_size 3
        lowerCAmelCase__   :				Dict 			=				2
        lowerCAmelCase__   :				Tuple 			=				inputs["input_ids"].shape[-1] // 2
        lowerCAmelCase__   :				Optional[Any] 			=				inputs["input_ids"][:max_batch_size, :sequence_length]

        lowerCAmelCase__   :				List[str] 			=				jnp.ones_like(snake_case   )
        lowerCAmelCase__   :				Tuple 			=				attention_mask[:max_batch_size, :sequence_length]

        # generate max 5 tokens
        lowerCAmelCase__   :				Optional[Any] 			=				input_ids.shape[-1] + 5
        if config.eos_token_id is not None and config.pad_token_id is None:
         # hack to allow generate for models such as GPT2 as is done in `generate()`
         lowerCAmelCase__   :				Tuple 			=				config.eos_token_id
        return config, input_ids, attention_mask, max_length






       @is_pt_flax_cross_test
       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				str 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Union[str, Any] 			=				False
        lowerCAmelCase__   :				Union[str, Any] 			=				max_length
        lowerCAmelCase__   :				List[str] 			=				0

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Dict 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Optional[Any] 			=				model_class.__name__[4:]  # Skip the "Flax" at the beginning
         lowerCAmelCase__   :				List[str] 			=				getattr(snake_case			,					snake_case   )
         lowerCAmelCase__   :				Any 			=				pt_model_class(snake_case   ).eval()
         lowerCAmelCase__   :				int 			=				load_flax_weights_in_pytorch_model(snake_case			,					flax_model.params   )

         lowerCAmelCase__   :				List[Any] 			=				flax_model.generate(snake_case   ).sequences
         lowerCAmelCase__   :				str 			=				pt_model.generate(torch.tensor(snake_case			,					dtype=torch.long   )   )

         if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
          lowerCAmelCase__   :				Any 			=				flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

         self.assertListEqual(pt_generation_outputs.numpy().tolist()			,					flax_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Optional[int] 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Union[str, Any] 			=				False
        lowerCAmelCase__   :				List[str] 			=				max_length

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				int 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Optional[Any] 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Tuple 			=				jit(model.generate   )
         lowerCAmelCase__   :				Union[str, Any] 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Dict 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Any 			=				True
        lowerCAmelCase__   :				int 			=				max_length

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				List[Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				List[Any] 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				str 			=				jit(model.generate   )
         lowerCAmelCase__   :				Dict 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Dict 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Union[str, Any] 			=				False
        lowerCAmelCase__   :				List[Any] 			=				max_length
        lowerCAmelCase__   :				Union[str, Any] 			=				2

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Optional[Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Dict 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				List[Any] 			=				jit(model.generate   )
         lowerCAmelCase__   :				Dict 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Any 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Optional[int] 			=				False
        lowerCAmelCase__   :				List[str] 			=				max_length
        lowerCAmelCase__   :				int 			=				2
        lowerCAmelCase__   :				Optional[Any] 			=				2

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				List[Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Any 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[0]			,					input_ids.shape[0] * config.num_return_sequences   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				List[Any] 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				int 			=				True
        lowerCAmelCase__   :				str 			=				max_length
        lowerCAmelCase__   :				int 			=				0.8
        lowerCAmelCase__   :				Tuple 			=				10
        lowerCAmelCase__   :				str 			=				0.3
        lowerCAmelCase__   :				Union[str, Any] 			=				1
        lowerCAmelCase__   :				Dict 			=				8
        lowerCAmelCase__   :				Optional[Any] 			=				9

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Union[str, Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Union[str, Any] 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Tuple 			=				jit(model.generate   )
         lowerCAmelCase__   :				Optional[Any] 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Tuple 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Union[str, Any] 			=				max_length
        lowerCAmelCase__   :				List[str] 			=				1
        lowerCAmelCase__   :				List[Any] 			=				8
        lowerCAmelCase__   :				Optional[int] 			=				9

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Union[str, Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Union[str, Any] 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Any 			=				jit(model.generate   )
         lowerCAmelCase__   :				List[Any] 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Any 			=				self._get_input_ids_and_config()
        lowerCAmelCase__   :				Tuple 			=				max_length
        lowerCAmelCase__   :				str 			=				2
        lowerCAmelCase__   :				List[str] 			=				1
        lowerCAmelCase__   :				Tuple 			=				8
        lowerCAmelCase__   :				Tuple 			=				9

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Tuple 			=				model_class(snake_case   )

         lowerCAmelCase__   :				int 			=				model.generate(snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Optional[Any] 			=				jit(model.generate   )
         lowerCAmelCase__   :				List[str] 			=				jit_generate(snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Tuple 			=				self._get_input_ids_and_config()

        # pad attention mask on the left
        lowerCAmelCase__   :				Optional[int] 			=				attention_mask.at[(0, 0)].set(0   )

        lowerCAmelCase__   :				str 			=				False
        lowerCAmelCase__   :				int 			=				max_length

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Dict 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Optional[Any] 			=				model.generate(snake_case			,					attention_mask=snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Any 			=				jit(model.generate   )
         lowerCAmelCase__   :				List[Any] 			=				jit_generate(snake_case			,					attention_mask=snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Tuple 			=				self._get_input_ids_and_config()

        # pad attention mask on the left
        lowerCAmelCase__   :				Union[str, Any] 			=				attention_mask.at[(0, 0)].set(0   )

        lowerCAmelCase__   :				List[Any] 			=				True
        lowerCAmelCase__   :				int 			=				max_length

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				List[str] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Tuple 			=				model.generate(snake_case			,					attention_mask=snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Tuple 			=				jit(model.generate   )
         lowerCAmelCase__   :				List[str] 			=				jit_generate(snake_case			,					attention_mask=snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Union[str, Any] 			=				self._get_input_ids_and_config()

        # pad attention mask on the left
        lowerCAmelCase__   :				Optional[Any] 			=				attention_mask.at[(0, 0)].set(0   )

        lowerCAmelCase__   :				List[str] 			=				2
        lowerCAmelCase__   :				str 			=				max_length

        for model_class in self.all_generative_model_classes:
         lowerCAmelCase__   :				Union[str, Any] 			=				model_class(snake_case   )

         lowerCAmelCase__   :				Optional[Any] 			=				model.generate(snake_case			,					attention_mask=snake_case   ).sequences
         self.assertEqual(generation_outputs.shape[-1]			,					snake_case   )

         lowerCAmelCase__   :				Any 			=				jit(model.generate   )
         lowerCAmelCase__   :				Tuple 			=				jit_generate(snake_case			,					attention_mask=snake_case   ).sequences

         self.assertListEqual(generation_outputs.tolist()			,					jit_generation_outputs.tolist()   )






@require_flax
class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				str 			=				AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"   )
        lowerCAmelCase__   :				List[Any] 			=				FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"   )

        lowerCAmelCase__   :				List[str] 			=				"Hello world"
        lowerCAmelCase__   :				str 			=				tokenizer(snake_case			,					return_tensors="np"   ).input_ids

        # typos are quickly detected (the correct argument is `do_sample`)
        with self.assertRaisesRegex(snake_case			,					"do_samples"   ):
         model.generate(snake_case			,					do_samples=snake_case   )

        # arbitrary arguments that will not be used anywhere are also not accepted
        with self.assertRaisesRegex(snake_case			,					"foo"   ):
         lowerCAmelCase__   :				List[Any] 			=				{"foo": "bar"}
         model.generate(snake_case			,					**snake_case   )









"""simple docstring"""

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_UpperCamelCase		     =    logging.get_logger(__name__)

_UpperCamelCase		     =    {
    """microsoft/focalnet-tiny""": """https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json""",
}






class 			__a					(						__magic_name__       ,	__magic_name__       ):





       """simple docstring"""


       __UpperCamelCase					:      Optional[Any]    =						'focalnet'






       def __init__(   self			,					snake_case=224			,					snake_case=4			,					snake_case=3			,					snake_case=96			,					snake_case=False			,					snake_case=[192, 384, 768, 768]			,					snake_case=[2, 2, 6, 2]			,					snake_case=[2, 2, 2, 2]			,					snake_case=[3, 3, 3, 3]			,					snake_case="gelu"			,					snake_case=4.0			,					snake_case=0.0			,					snake_case=0.1			,					snake_case=False			,					snake_case=1e-4			,					snake_case=False			,					snake_case=False			,					snake_case=False			,					snake_case=0.02			,					snake_case=1e-5			,					snake_case=32			,					snake_case=None			,					snake_case=None			,					**snake_case			,					):





        """simple docstring"""

        super().__init__(**snake_case   )

        lowerCAmelCase__   :				Optional[Any] 			=				image_size
        lowerCAmelCase__   :				Union[str, Any] 			=				patch_size
        lowerCAmelCase__   :				List[str] 			=				num_channels
        lowerCAmelCase__   :				List[str] 			=				embed_dim
        lowerCAmelCase__   :				List[Any] 			=				use_conv_embed
        lowerCAmelCase__   :				List[str] 			=				hidden_sizes
        lowerCAmelCase__   :				List[Any] 			=				depths
        lowerCAmelCase__   :				Union[str, Any] 			=				focal_levels
        lowerCAmelCase__   :				Union[str, Any] 			=				focal_windows
        lowerCAmelCase__   :				int 			=				hidden_act
        lowerCAmelCase__   :				List[Any] 			=				mlp_ratio
        lowerCAmelCase__   :				str 			=				hidden_dropout_prob
        lowerCAmelCase__   :				Union[str, Any] 			=				drop_path_rate
        lowerCAmelCase__   :				Tuple 			=				use_layerscale
        lowerCAmelCase__   :				Tuple 			=				layerscale_value
        lowerCAmelCase__   :				str 			=				use_post_layernorm
        lowerCAmelCase__   :				str 			=				use_post_layernorm_in_modulation
        lowerCAmelCase__   :				Union[str, Any] 			=				normalize_modulator
        lowerCAmelCase__   :				Optional[int] 			=				initializer_range
        lowerCAmelCase__   :				List[Any] 			=				layer_norm_eps
        lowerCAmelCase__   :				str 			=				encoder_stride
        lowerCAmelCase__   :				Union[str, Any] 			=				["stem"] + [F"""stage{idx}""" for idx in range(1			,					len(self.depths   ) + 1   )]
        lowerCAmelCase__			,						lowerCAmelCase__   :				Union[str, Any] 			=				get_aligned_output_features_output_indices(
            out_features=snake_case			,					out_indices=snake_case			,					stage_names=self.stage_names   )






