
'''simple docstring'''
from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
     from PIL import Image

     from transformers import ConvNextImageProcessor







class 			__A					:
 def __init__( self      :     Optional[int]					, lowerCamelCase      :     Optional[int]					, lowerCamelCase      :     Any=13					, lowerCamelCase      :     Union[str, Any]=32					, lowerCamelCase      :     Union[str, Any]=3					, lowerCamelCase      :     List[str]=4					, lowerCamelCase      :     str=[10, 20, 30, 40]					, lowerCamelCase      :     Tuple=[2, 2, 3, 2]					, lowerCamelCase      :     Union[str, Any]=True					, lowerCamelCase      :     Optional[int]=True					, lowerCamelCase      :     List[str]=37					, lowerCamelCase      :     int="gelu"					, lowerCamelCase      :     Optional[int]=10					, lowerCamelCase      :     Union[str, Any]=0.02					, lowerCamelCase      :     str=3					, lowerCamelCase      :     Union[str, Any]=None					, ):



   """simple docstring"""





   __A		:  Any       				=		parent
   __A		:  List[Any]       				=		batch_size
   __A		:  Optional[int]       				=		image_size
   __A		:  Optional[Any]       				=		num_channels
   __A		:  Optional[Any]       				=		num_stages
   __A		:  Any       				=		hidden_sizes
   __A		:  List[str]       				=		depths
   __A		:  List[str]       				=		is_training
   __A		:  int       				=		use_labels
   __A		:  Union[str, Any]       				=		intermediate_size
   __A		:  Optional[int]       				=		hidden_act
   __A		:  List[str]       				=		type_sequence_label_size
   __A		:  int       				=		initializer_range
   __A		:  Union[str, Any]       				=		scope
 def 		lowercase_( self      :     List[str]							):



   """simple docstring"""





   __A		:  List[Any]       				=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

   __A		:  Union[str, Any]       				=		None
   if self.use_labels:
     __A		:  Any       				=		ids_tensor([self.batch_size]					, self.type_sequence_label_size							)

   __A		:  Tuple       				=		self.get_config()

   return config, pixel_values, labels
 def 		lowercase_( self      :     Tuple							):



   """simple docstring"""





   return ConvNextConfig(
       num_channels=self.num_channels					, hidden_sizes=self.hidden_sizes					, depths=self.depths					, num_stages=self.num_stages					, hidden_act=self.hidden_act					, is_decoder=a__					, initializer_range=self.initializer_range					, )
 def 		lowercase_( self      :     Optional[int]					, lowerCamelCase      :     int					, lowerCamelCase      :     Union[str, Any]					, lowerCamelCase      :     int							):



   """simple docstring"""





   __A		:  Dict       				=		TFConvNextModel(config=a__							)
   __A		:  List[Any]       				=		model(a__					, training=a__							)
   # expected last hidden states: B, C, H // 32, W // 32
   self.parent.assertEqual(
       result.last_hidden_state.shape					, (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)					, )
 def 		lowercase_( self      :     Union[str, Any]					, lowerCamelCase      :     List[str]					, lowerCamelCase      :     Optional[int]					, lowerCamelCase      :     str							):



   """simple docstring"""





   __A		:  Optional[Any]       				=		self.type_sequence_label_size
   __A		:  Optional[Any]       				=		TFConvNextForImageClassification(a__							)
   __A		:  Optional[Any]       				=		model(a__					, labels=a__					, training=a__							)
   self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.type_sequence_label_size)							)






 def 		lowercase_( self      :     int							):



   """simple docstring"""





   __A		:  Optional[Any]       				=		self.prepare_config_and_inputs()
   __A		:  Dict       				=		config_and_inputs
   __A		:  List[Any]       				=		{'''pixel_values''': pixel_values}
   return config, inputs_dict







@require_tf
class 			__A					(			_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    unittest.TestCase      ):
 lowerCamelCase												=(TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
 lowerCamelCase												=(
     {'''feature-extraction''': TFConvNextModel, '''image-classification''': TFConvNextForImageClassification}
     if is_tf_available()
     else {}
 )

 lowerCamelCase												=False
 lowerCamelCase												=False
 lowerCamelCase												=False
 lowerCamelCase												=False
 lowerCamelCase												=False
 def 		lowercase_( self      :     int							):



   """simple docstring"""





   __A		:  str       				=		TFConvNextModelTester(self							)
   __A		:  List[Any]       				=		ConfigTester(
       self					, config_class=a__					, has_text_modality=a__					, hidden_size=37					, )
 @unittest.skip(reason="""ConvNext does not use inputs_embeds"""							)
 def 		lowercase_( self      :     str							):



   """simple docstring"""





   pass
 @unittest.skipIf(
     not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""							)							) == 0					, reason="""TF does not support backprop for grouped convolutions on CPU."""					, )
 @slow
 def 		lowercase_( self      :     Union[str, Any]							):



   """simple docstring"""





   super().test_keras_fit()
 @unittest.skip(reason="""ConvNext does not support input and output embeddings"""							)
 def 		lowercase_( self      :     List[str]							):



   """simple docstring"""





   pass
 def 		lowercase_( self      :     int							):



   """simple docstring"""





   __A		:  Any       				=		self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     __A		:  Optional[int]       				=		model_class(a__							)
     __A		:  List[str]       				=		inspect.signature(model.call							)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     __A		:  List[str]       				=		[*signature.parameters.keys()]

     __A		:  Dict       				=		['''pixel_values''']
     self.assertListEqual(arg_names[:1]					, a__							)
 def 		lowercase_( self      :     str							):



   """simple docstring"""





   __A		:  Optional[Any]       				=		self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a__							)
 @unittest.skipIf(
     not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""							)							) == 0					, reason="""TF does not support backprop for grouped convolutions on CPU."""					, )
 def 		lowercase_( self      :     int							):



   """simple docstring"""





   super().test_dataset_conversion()
 def 		lowercase_( self      :     Optional[int]							):



   """simple docstring"""





   def check_hidden_states_output(lowerCamelCase      :     Any					, lowerCamelCase      :     Tuple					, lowerCamelCase      :     List[Any]							):
     __A		:  Union[str, Any]       				=		model_class(a__							)

     __A		:  str       				=		model(**self._prepare_for_class(a__					, a__							)							)
     __A		:  Any       				=		outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

     __A		:  Optional[int]       				=		self.model_tester.num_stages
     self.assertEqual(len(a__							)					, expected_num_stages + 1							)

     # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
     self.assertListEqual(
         list(hidden_states[0].shape[-2:]							)					, [self.model_tester.image_size // 4, self.model_tester.image_size // 4]					, )

   __A		:  Dict       				=		self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     __A		:  Optional[int]       				=		True
     check_hidden_states_output(a__					, a__					, a__							)

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

     check_hidden_states_output(a__					, a__					, a__							)
 def 		lowercase_( self      :     Union[str, Any]							):



   """simple docstring"""





   __A		:  Any       				=		self.model_tester.prepare_config_and_inputs_for_common()

   def check_equivalence(lowerCamelCase      :     Optional[int]					, lowerCamelCase      :     int					, lowerCamelCase      :     Dict					, lowerCamelCase      :     Tuple={}							):
     __A		:  List[Any]       				=		model(a__					, return_dict=a__					, **a__							)
     __A		:  List[str]       				=		model(a__					, return_dict=a__					, **a__							).to_tuple()

     def recursive_check(lowerCamelCase      :     Union[str, Any]					, lowerCamelCase      :     Union[str, Any]							):
       if isinstance(a__					, (List, Tuple)							):
         for tuple_iterable_value, dict_iterable_value in zip(a__					, a__							):
           recursive_check(a__					, a__							)
       elif tuple_object is None:
         return
       else:
         self.assertTrue(
             all(tf.equal(a__					, a__							)							)					, msg=(
                 """Tuple and dict output are not equal. Difference:"""
                 f" {tf.math.reduce_max(tf.abs(tuple_object - dict_object							)							)}"
             )					, )

       recursive_check(a__					, a__							)

   for model_class in self.all_model_classes:
     __A		:  List[str]       				=		model_class(a__							)

     __A		:  int       				=		self._prepare_for_class(a__					, a__							)
     __A		:  Any       				=		self._prepare_for_class(a__					, a__							)
     check_equivalence(a__					, a__					, a__							)

     __A		:  Optional[int]       				=		self._prepare_for_class(a__					, a__					, return_labels=a__							)
     __A		:  Union[str, Any]       				=		self._prepare_for_class(a__					, a__					, return_labels=a__							)
     check_equivalence(a__					, a__					, a__							)

     __A		:  Optional[Any]       				=		self._prepare_for_class(a__					, a__							)
     __A		:  Optional[Any]       				=		self._prepare_for_class(a__					, a__							)
     check_equivalence(a__					, a__					, a__					, {"""output_hidden_states""": True}							)

     __A		:  Optional[Any]       				=		self._prepare_for_class(a__					, a__					, return_labels=a__							)
     __A		:  List[Any]       				=		self._prepare_for_class(a__					, a__					, return_labels=a__							)
     check_equivalence(a__					, a__					, a__					, {"""output_hidden_states""": True}							)
 def 		lowercase_( self      :     int							):



   """simple docstring"""





   __A		:  Union[str, Any]       				=		self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*a__							)






 @slow
 def 		lowercase_( self      :     Optional[int]							):



   """simple docstring"""





   __A		:  Tuple       				=		TFConvNextModel.from_pretrained("""facebook/convnext-tiny-224"""							)
   self.assertIsNotNone(a__							)







def 							A_  (    )					->							int:





  """simple docstring"""

  __A		:  Any       				=		Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
  return image







@require_tf
@require_vision
class 			__A					(			unittest.TestCase      ):
 @cached_property
 def 		lowercase_( self      :     int							):



   """simple docstring"""





   return ConvNextImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""							) if is_vision_available() else None






 @slow
 def 		lowercase_( self      :     Any							):



   """simple docstring"""





   __A		:  Optional[int]       				=		TFConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""							)

   __A		:  Optional[Any]       				=		self.default_image_processor
   __A		:  Tuple       				=		prepare_img()
   __A		:  Dict       				=		image_processor(images=a__					, return_tensors="""tf"""							)

   # forward pass
   __A		:  Union[str, Any]       				=		model(**a__							)

   # verify the logits
   __A		:  List[Any]       				=		tf.TensorShape((1, 10_00)							)
   self.assertEqual(outputs.logits.shape					, a__							)

   __A		:  List[Any]       				=		tf.constant([-0.0_260, -0.4_739, 0.1_911]							)

   tf.debugging.assert_near(outputs.logits[0, :3]					, a__					, atol=1e-4							)







import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp
  import numpy as np

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
  import torch

@require_flax
class   lowerCamelCase						:

       def __init__(      self						,    a__						,    a__=14						,    a__=7						,    a__=True						,    a__=True						,    a__=True						,    a__=99						,    a__=32						,    a__=5						,    a__=4						,    a__=37						,    a__="gelu"						,    a__=0.1						,    a__=0.1						,    a__=5_12						,    a__=0.02						,    a__=None						,    ):






        '''simple docstring'''






        snake_case :			Any							=   parent
        snake_case :			str							=   batch_size
        snake_case :			str							=   seq_length
        snake_case :			str							=   is_training
        snake_case :			List[str]							=   use_input_mask
        snake_case :			str							=   use_labels
        snake_case :			Union[str, Any]							=   vocab_size
        snake_case :			Dict							=   d_model
        snake_case :			List[str]							=   num_hidden_layers
        snake_case :			Optional[int]							=   num_attention_heads
        snake_case :			Optional[Any]							=   ffn_dim
        snake_case :			Dict							=   activation_function
        snake_case :			List[str]							=   activation_dropout
        snake_case :			Any							=   attention_dropout
        snake_case :			str							=   max_position_embeddings
        snake_case :			Optional[int]							=   initializer_range
        snake_case :			Union[str, Any]							=   None
        snake_case :			int							=   0
        snake_case :			int							=   2
        snake_case :			Union[str, Any]							=   1

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   np.clip(ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size       )						,    3						,    self.vocab_size       )

        snake_case :			Union[str, Any]							=   None
        if self.use_input_mask:
         snake_case :			Tuple							=   random_attention_mask([self.batch_size, self.seq_length]       )

        snake_case :			Optional[Any]							=   XGLMConfig(
            vocab_size=self.vocab_size						,    d_model=self.hidden_size						,    num_layers=self.num_hidden_layers						,    attention_heads=self.num_attention_heads						,    ffn_dim=self.ffn_dim						,    activation_function=self.activation_function						,    activation_dropout=self.activation_dropout						,    attention_dropout=self.attention_dropout						,    max_position_embeddings=self.max_position_embeddings						,    initializer_range=self.initializer_range						,    use_cache=a__						,    bos_token_id=self.bos_token_id						,    eos_token_id=self.eos_token_id						,    pad_token_id=self.pad_token_id						,    )

        return (config, input_ids, input_mask)

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Optional[int]							=   self.prepare_config_and_inputs()
        snake_case	,       snake_case	,       snake_case :			List[str]							=   config_and_inputs
        snake_case :			List[str]							=   {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
        return config, inputs_dict

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case	,       snake_case :			List[str]							=   self.prepare_config_and_inputs()

        snake_case :			Tuple							=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
        snake_case :			Union[str, Any]							=   ids_tensor([self.batch_size, self.seq_length]						,    vocab_size=2       )

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )

       def 					UpperCAmelCase					(      self						,    a__						,    a__						,    a__						,    a__       ):






        '''simple docstring'''






        snake_case :			Optional[Any]							=   20
        snake_case :			List[Any]							=   model_class_name(a__       )

        snake_case :			List[Any]							=   model.init_cache(input_ids.shape[0]						,    a__       )
        snake_case :			Optional[int]							=   jnp.ones((input_ids.shape[0], max_decoder_length)						,    dtype='''i4'''       )

        snake_case :			Optional[Any]							=   jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)       )
        snake_case :			Tuple							=   model(
            input_ids[:, :-1]						,    attention_mask=a__						,    past_key_values=a__						,    position_ids=a__						,    )

        snake_case :			Dict							=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype='''i4'''       )
        snake_case :			str							=   model(
            input_ids[:, -1:]						,    attention_mask=a__						,    past_key_values=outputs_cache.past_key_values						,    position_ids=a__						,    )

        snake_case :			Dict							=   model(a__       )

        snake_case :			Union[str, Any]							=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
        self.parent.assertTrue(diff < 1e-3						,    msg=f'''Max diff is {diff}'''       )



       def 					UpperCAmelCase					(      self						,    a__						,    a__						,    a__						,    a__       ):






        '''simple docstring'''






        snake_case :			Optional[int]							=   20
        snake_case :			Dict							=   model_class_name(a__       )

        snake_case :			Dict							=   jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])       )]						,    axis=-1						,    )

        snake_case :			List[str]							=   model.init_cache(input_ids.shape[0]						,    a__       )
        snake_case :			Optional[int]							=   jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)       )

        snake_case :			Optional[int]							=   model(
            input_ids[:, :-1]						,    attention_mask=a__						,    past_key_values=a__						,    position_ids=a__						,    )
        snake_case :			int							=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype='''i4'''       )
        snake_case :			Optional[int]							=   model(
            input_ids[:, -1:]						,    past_key_values=outputs_cache.past_key_values						,    attention_mask=a__						,    position_ids=a__						,    )

        snake_case :			Optional[Any]							=   model(a__						,    attention_mask=a__       )
        snake_case :			Optional[int]							=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
        self.parent.assertTrue(diff < 1e-3						,    msg=f'''Max diff is {diff}'''       )







@require_sentencepiece
@require_flax
class   lowerCamelCase						(	SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      unittest.TestCase  ):
       _UpperCamelCase     		=   (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
       _UpperCamelCase     		=   (FlaxXGLMForCausalLM,) if is_flax_available() else ()

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   FlaxXGLMModelTester(self       )

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case	,       snake_case	,       snake_case :			str							=   self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward(a__						,    a__						,    a__						,    a__       )

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case	,       snake_case	,       snake_case :			int							=   self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward_with_attn_mask(
             a__						,    a__						,    a__						,    a__       )

       @slow
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   XGLMTokenizer.from_pretrained('''XGLM'''						,    padding_side='''left'''       )
        snake_case :			Optional[int]							=   tokenizer(['''Hello this is a long string''', '''Hey''']						,    return_tensors='''np'''						,    padding=a__						,    truncation=a__       )

        snake_case :			List[Any]							=   FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''       )
        snake_case :			Optional[int]							=   1
        snake_case :			str							=   False

        snake_case :			str							=   jax.jit(model.generate       )

        snake_case :			Dict							=   jit_generate(inputs['''input_ids''']						,    attention_mask=inputs['''attention_mask''']       ).sequences

        snake_case :			Optional[Any]							=   tokenizer.batch_decode(a__						,    skip_special_tokens=a__       )

        snake_case :			Tuple							=   [
            '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
            '''Hey, I\'m a newbie to the forum and I\'''',
        ]

        self.assertListEqual(a__						,    a__       )

       @is_pt_flax_cross_test
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case :			Union[str, Any]							=   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__       ):
          # prepare inputs
          snake_case :			Union[str, Any]							=   self._prepare_for_class(a__						,    a__       )
          snake_case :			int							=   {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          snake_case :			Any							=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          snake_case :			Optional[Any]							=   getattr(a__						,    a__       )

          snake_case	,       snake_case :			Optional[Any]							=   pt_inputs['''input_ids'''].shape
          snake_case :			Tuple							=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)       )
          for batch_idx, start_index in enumerate(a__       ):
           snake_case :			Optional[Any]							=   0
           snake_case :			Optional[int]							=   1
           snake_case :			Optional[Any]							=   0
           snake_case :			Optional[Any]							=   1
          snake_case :			Dict							=   pt_model_class(a__       ).eval()
          # Flax models don't use the `use_cache` option and cache is not returned as a default.
          # So we disable `use_cache` here for PyTorch model.
          snake_case :			Any							=   False
          snake_case :			str							=   model_class(a__						,    dtype=jnp.floataa       )
          snake_case :			Optional[int]							=   convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,    a__       )
          snake_case :			Optional[int]							=   fx_state

          with torch.no_grad():
           snake_case :			int							=   pt_model(**a__       ).to_tuple()

          snake_case :			Optional[int]							=   fx_model(**a__       ).to_tuple()
          self.assertEqual(len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

          with tempfile.TemporaryDirectory() as tmpdirname:
           pt_model.save_pretrained(a__       )
           snake_case :			Union[str, Any]							=   model_class.from_pretrained(a__						,    from_pt=a__       )

          snake_case :			List[Any]							=   fx_model_loaded(**a__       ).to_tuple()
          self.assertEqual(
              len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output_loaded, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output_loaded[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

       @is_pt_flax_cross_test
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case :			List[Any]							=   self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__       ):
          # prepare inputs
          snake_case :			str							=   self._prepare_for_class(a__						,    a__       )
          snake_case :			List[Any]							=   {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          snake_case :			Any							=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          snake_case :			Union[str, Any]							=   getattr(a__						,    a__       )

          snake_case :			List[str]							=   pt_model_class(a__       ).eval()
          snake_case :			Tuple							=   False
          snake_case :			Optional[int]							=   model_class(a__						,    dtype=jnp.floataa       )

          snake_case :			Any							=   load_flax_weights_in_pytorch_model(a__						,    fx_model.params       )
          snake_case	,       snake_case :			Union[str, Any]							=   pt_inputs['''input_ids'''].shape
          snake_case :			List[Any]							=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)       )
          for batch_idx, start_index in enumerate(a__       ):
           snake_case :			Union[str, Any]							=   0
           snake_case :			Tuple							=   1
           snake_case :			int							=   0
           snake_case :			Optional[int]							=   1

          # make sure weights are tied in PyTorch
          pt_model.tie_weights()

          with torch.no_grad():
           snake_case :			Optional[int]							=   pt_model(**a__       ).to_tuple()

          snake_case :			Any							=   fx_model(**a__       ).to_tuple()
          self.assertEqual(len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

          with tempfile.TemporaryDirectory() as tmpdirname:
           fx_model.save_pretrained(a__       )
           snake_case :			int							=   pt_model_class.from_pretrained(a__						,    from_flax=a__       )

          with torch.no_grad():
           snake_case :			int							=   pt_model_loaded(**a__       ).to_tuple()

          self.assertEqual(
              len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )



       @slow
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case :			str							=   model_class_name.from_pretrained('''facebook/xglm-564M'''       )
         snake_case :			List[Any]							=   model(np.ones((1, 1)       )       )
         self.assertIsNotNone(a__       )






