




'''simple docstring'''

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






						def __init__(   self						:						Tuple					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Tuple=1_4					,   __UpperCamelCase						:						List[Any]=7					,   __UpperCamelCase						:						List[str]=True					,   __UpperCamelCase						:						Tuple=True					,   __UpperCamelCase						:						Any=True					,   __UpperCamelCase						:						str=9_9					,   __UpperCamelCase						:						str=3_2					,   __UpperCamelCase						:						Optional[Any]=5					,   __UpperCamelCase						:						int=4					,   __UpperCamelCase						:						Union[str, Any]=3_7					,   __UpperCamelCase						:						List[Any]="gelu"					,   __UpperCamelCase						:						List[Any]=0.1					,   __UpperCamelCase						:						Tuple=0.1					,   __UpperCamelCase						:						int=5_1_2					,   __UpperCamelCase						:						Tuple=0.02					,   __UpperCamelCase						:						Dict=None					,   ):



									'''simple docstring'''






									snake_case__          =						parent
									snake_case__          =						batch_size
									snake_case__          =						seq_length
									snake_case__          =						is_training
									snake_case__          =						use_input_mask
									snake_case__          =						use_labels
									snake_case__          =						vocab_size
									snake_case__          =						d_model
									snake_case__          =						num_hidden_layers
									snake_case__          =						num_attention_heads
									snake_case__          =						ffn_dim
									snake_case__          =						activation_function
									snake_case__          =						activation_dropout
									snake_case__          =						attention_dropout
									snake_case__          =						max_position_embeddings
									snake_case__          =						initializer_range
									snake_case__          =						None
									snake_case__          =						0
									snake_case__          =						2
									snake_case__          =						1






						def       __lowerCAmelCase(   self						:						List[str]       ):



									'''simple docstring'''






									snake_case__          =						np.clip(ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size       )					,   3					,   self.vocab_size       )

									snake_case__          =						None
									if self.use_input_mask:
												snake_case__          =						random_attention_mask([self.batch_size, self.seq_length]       )

									snake_case__          =						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       __lowerCAmelCase(   self						:						str       ):



									'''simple docstring'''






									snake_case__          =						self.prepare_config_and_inputs()
									snake_case__          =						config_and_inputs
									snake_case__          =						{"input_ids": input_ids, "attention_mask": attention_mask}
									return config, inputs_dict






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									snake_case__          =						self.prepare_config_and_inputs()

									snake_case__          =						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
									snake_case__          =						ids_tensor([self.batch_size, self.seq_length]					,   vocab_size=2       )

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






						def       __lowerCAmelCase(   self						:						int					,   __UpperCamelCase						:						Optional[int]					,   __UpperCamelCase						:						Dict					,   __UpperCamelCase						:						Tuple					,   __UpperCamelCase						:						Dict       ):



									'''simple docstring'''






									snake_case__          =						2_0
									snake_case__          =						model_class_name(_a       )

									snake_case__          =						model.init_cache(input_ids.shape[0]					,   _a       )
									snake_case__          =						jnp.ones((input_ids.shape[0], max_decoder_length)					,   dtype="""i4"""       )

									snake_case__          =						jnp.broadcast_to(
									    jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,   (input_ids.shape[0], input_ids.shape[-1] - 1)       )
									snake_case__          =						model(
									    input_ids[:, :-1]					,   attention_mask=_a					,   past_key_values=_a					,   position_ids=_a					,   )

									snake_case__          =						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,   dtype="""i4"""       )
									snake_case__          =						model(
									    input_ids[:, -1:]					,   attention_mask=_a					,   past_key_values=outputs_cache.past_key_values					,   position_ids=_a					,   )

									snake_case__          =						model(_a       )

									snake_case__          =						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       __lowerCAmelCase(   self						:						str					,   __UpperCamelCase						:						List[Any]					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						Optional[Any]					,   __UpperCamelCase						:						List[Any]       ):



									'''simple docstring'''






									snake_case__          =						2_0
									snake_case__          =						model_class_name(_a       )

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

									snake_case__          =						model.init_cache(input_ids.shape[0]					,   _a       )
									snake_case__          =						jnp.broadcast_to(
									    jnp.arange(input_ids.shape[-1] - 1       )[None, :]					,   (input_ids.shape[0], input_ids.shape[-1] - 1)       )

									snake_case__          =						model(
									    input_ids[:, :-1]					,   attention_mask=_a					,   past_key_values=_a					,   position_ids=_a					,   )
									snake_case__          =						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,   dtype="""i4"""       )
									snake_case__          =						model(
									    input_ids[:, -1:]					,   past_key_values=outputs_cache.past_key_values					,   attention_mask=_a					,   position_ids=_a					,   )

									snake_case__          =						model(_a					,   attention_mask=_a       )
									snake_case__          =						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 __magic_name__(						a_    ,		a_    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Tuple    =  (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
						UpperCAmelCase_		:					int    =  (FlaxXGLMForCausalLM,) if is_flax_available() else ()






						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									snake_case__          =						FlaxXGLMModelTester(self       )






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									for model_class_name in self.all_model_classes:
												snake_case__          =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.check_use_cache_forward(_a					,   _a					,   _a					,   _a       )






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									for model_class_name in self.all_model_classes:
												snake_case__          =						self.model_tester.prepare_config_and_inputs()
												self.model_tester.check_use_cache_forward_with_attn_mask(
												    _a					,   _a					,   _a					,   _a       )






						@slow
						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									snake_case__          =						XGLMTokenizer.from_pretrained("""XGLM"""					,   padding_side="""left"""       )
									snake_case__          =						tokenizer(["""Hello this is a long string""", """Hey"""]					,   return_tensors="""np"""					,   padding=_a					,   truncation=_a       )

									snake_case__          =						FlaxXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""       )
									snake_case__          =						1
									snake_case__          =						False

									snake_case__          =						jax.jit(model.generate       )

									snake_case__          =						jit_generate(inputs["""input_ids"""]					,   attention_mask=inputs["""attention_mask"""]       ).sequences

									snake_case__          =						tokenizer.batch_decode(_a					,   skip_special_tokens=_a       )

									snake_case__          =						[
									    "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       __lowerCAmelCase(   self						:						Union[str, Any]       ):



									'''simple docstring'''






									snake_case__          =						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__          =						self._prepare_for_class(_a					,   _a       )
															snake_case__          =						{k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

															# load corresponding PyTorch class
															snake_case__          =						model_class.__name__[4:]  # Skip the "Flax" at the beginning
															snake_case__          =						getattr(_a					,   _a       )

															snake_case__          =						pt_inputs["input_ids"].shape
															snake_case__          =						np.random.randint(0					,   seq_length - 1					,   size=(batch_size,)       )
															for batch_idx, start_index in enumerate(_a       ):
																		snake_case__          =						0
																		snake_case__          =						1
																		snake_case__          =						0
																		snake_case__          =						1
															snake_case__          =						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__          =						False
															snake_case__          =						model_class(_a					,   dtype=jnp.floataa       )
															snake_case__          =						convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,   _a       )
															snake_case__          =						fx_state

															with torch.no_grad():
																		snake_case__          =						pt_model(**_a       ).to_tuple()

															snake_case__          =						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__          =						model_class.from_pretrained(_a					,   from_pt=_a       )

															snake_case__          =						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       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									snake_case__          =						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__          =						self._prepare_for_class(_a					,   _a       )
															snake_case__          =						{k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

															# load corresponding PyTorch class
															snake_case__          =						model_class.__name__[4:]  # Skip the "Flax" at the beginning
															snake_case__          =						getattr(_a					,   _a       )

															snake_case__          =						pt_model_class(_a       ).eval()
															snake_case__          =						False
															snake_case__          =						model_class(_a					,   dtype=jnp.floataa       )

															snake_case__          =						load_flax_weights_in_pytorch_model(_a					,   fx_model.params       )
															snake_case__          =						pt_inputs["input_ids"].shape
															snake_case__          =						np.random.randint(0					,   seq_length - 1					,   size=(batch_size,)       )
															for batch_idx, start_index in enumerate(_a       ):
																		snake_case__          =						0
																		snake_case__          =						1
																		snake_case__          =						0
																		snake_case__          =						1

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

															with torch.no_grad():
																		snake_case__          =						pt_model(**_a       ).to_tuple()

															snake_case__          =						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__          =						pt_model_class.from_pretrained(_a					,   from_flax=_a       )

															with torch.no_grad():
																		snake_case__          =						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       __lowerCAmelCase(   self						:						Any       ):



									'''simple docstring'''






									for model_class_name in self.all_model_classes:
												snake_case__          =						model_class_name.from_pretrained("""facebook/xglm-564M"""       )
												snake_case__          =						model(np.ones((1, 1)       )       )
												self.assertIsNotNone(_a       )



from __future__ import annotations

import sys







def 	A (		A__		,		A__			):
     # base cases
     if abs(A__			) == 1:
          return a, 0
     elif abs(A__			) == 1:
          return 0, b

     UpperCAmelCase	,			UpperCAmelCase    :					Any			=      a, b
     UpperCAmelCase	,			UpperCAmelCase    :					List[Any]			=      1, 0
     UpperCAmelCase	,			UpperCAmelCase    :					Union[str, Any]			=      0, 1

     while remainder != 0:
          UpperCAmelCase    :					Union[str, Any]			=      old_remainder // remainder
          UpperCAmelCase	,			UpperCAmelCase    :					Dict			=      remainder, old_remainder - quotient * remainder
          UpperCAmelCase	,			UpperCAmelCase    :					int			=      coeff_a, old_coeff_a - quotient * coeff_a
          UpperCAmelCase	,			UpperCAmelCase    :					Optional[Any]			=      coeff_b, old_coeff_b - quotient * coeff_b

     # sign correction for negative numbers
     if a < 0:
          UpperCAmelCase    :					Union[str, Any]			=      -old_coeff_a
     if b < 0:
          UpperCAmelCase    :					int			=      -old_coeff_b

     return old_coeff_a, old_coeff_b







def 	A (		):
     if len(sys.argv			) < 3:
          print("2 integer arguments required"			)
          return 1
     UpperCAmelCase    :					Union[str, Any]			=      int(sys.argv[1]			)
     UpperCAmelCase    :					Any			=      int(sys.argv[2]			)
     print(extended_euclidean_algorithm(A__		,		A__			)			)
     return 0


if __name__ == "__main__":
      raise SystemExit(main())



