



'''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      lowerCamelCase__     :






	'''simple docstring'''

	def __init__(  self		:  Dict   ,		__A		:  int   ,		__A		:  Tuple=14   ,		__A		:  int=7   ,		__A		:  str=True   ,		__A		:  Optional[int]=True   ,		__A		:  Any=True   ,		__A		:  Optional[int]=99   ,		__A		:  Tuple=32   ,		__A		:  Optional[int]=5   ,		__A		:  List[Any]=4   ,		__A		:  List[Any]=37   ,		__A		:  Dict="gelu"   ,		__A		:  Dict=0.1   ,		__A		:  Optional[int]=0.1   ,		__A		:  Tuple=512   ,		__A		:  Dict=0.0_2   ,		__A		:  Union[str, Any]=None   ,		)     ->			str:




								'''simple docstring'''





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




	def      lowercase__							(  self		:  Optional[int]					)     ->			Union[str, Any]:




								'''simple docstring'''





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

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

								lowerCAmelCase__  =   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      lowercase__							(  self		:  Optional[Any]					)     ->			Tuple:




								'''simple docstring'''





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




	def      lowercase__							(  self		:  int					)     ->			int:




								'''simple docstring'''





								lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.prepare_config_and_inputs()

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

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




	def      lowercase__							(  self		:  Dict   ,		__A		:  Dict   ,		__A		:  Dict   ,		__A		:  Tuple   ,		__A		:  List[Any]					)     ->			Union[str, Any]:




								'''simple docstring'''





								lowerCAmelCase__  =   20
								lowerCAmelCase__  =   model_class_name(__A					)

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

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

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

								lowerCAmelCase__  =   model(__A					)

								lowerCAmelCase__  =   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      lowercase__							(  self		:  Dict   ,		__A		:  List[Any]   ,		__A		:  Tuple   ,		__A		:  str   ,		__A		:  Any					)     ->			List[str]:




								'''simple docstring'''





								lowerCAmelCase__  =   20
								lowerCAmelCase__  =   model_class_name(__A					)

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

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

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

								lowerCAmelCase__  =   model(__A   ,		attention_mask=__A					)
								lowerCAmelCase__  =   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__     (     _A,     _A,     unittest.TestCase				):






	'''simple docstring'''

	A__		=	(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
	A__		=	(FlaxXGLMForCausalLM,) if is_flax_available() else ()




	def      lowercase__							(  self		:  Optional[int]					)     ->			List[str]:




								'''simple docstring'''





								lowerCAmelCase__  =   FlaxXGLMModelTester(self					)




	def      lowercase__							(  self		:  Tuple					)     ->			Union[str, Any]:




								'''simple docstring'''





								for model_class_name in self.all_model_classes:
															lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward(__A   ,		__A   ,		__A   ,		__A					)




	def      lowercase__							(  self		:  Any					)     ->			List[Any]:




								'''simple docstring'''





								for model_class_name in self.all_model_classes:
															lowerCAmelCase__  ,lowerCAmelCase__  ,lowerCAmelCase__  =   self.model_tester.prepare_config_and_inputs()
															self.model_tester.check_use_cache_forward_with_attn_mask(
															    __A   ,		__A   ,		__A   ,		__A					)




	@slow
	def      lowercase__							(  self		:  str					)     ->			Any:




								'''simple docstring'''





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

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

								lowerCAmelCase__  =   jax.jit(model.generate					)

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

								lowerCAmelCase__  =   tokenizer.batch_decode(__A   ,		skip_special_tokens=__A					)

								lowerCAmelCase__  =   [
								    """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      lowercase__							(  self		:  Tuple					)     ->			Dict:




								'''simple docstring'''





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

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

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

																						with torch.no_grad():
																													lowerCAmelCase__  =   pt_model(**__A					).to_tuple()

																						lowerCAmelCase__  =   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					)
																													lowerCAmelCase__  =   model_class.from_pretrained(__A   ,		from_pt=__A					)

																						lowerCAmelCase__  =   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      lowercase__							(  self		:  int					)     ->			Union[str, Any]:




								'''simple docstring'''





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

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

																						lowerCAmelCase__  =   pt_model_class(__A					).eval()
																						lowerCAmelCase__  =   False
																						lowerCAmelCase__  =   model_class(__A   ,		dtype=jnp.floataa					)

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

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

																						with torch.no_grad():
																													lowerCAmelCase__  =   pt_model(**__A					).to_tuple()

																						lowerCAmelCase__  =   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					)
																													lowerCAmelCase__  =   pt_model_class.from_pretrained(__A   ,		from_flax=__A					)

																						with torch.no_grad():
																													lowerCAmelCase__  =   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      lowercase__							(  self		:  Union[str, Any]					)     ->			Optional[int]:




								'''simple docstring'''





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










'''simple docstring'''







from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	bool						) ->							torch.Tensor:
							# The first operation in a checkpoint can't be in-place, but it's
							# nice to have in-place addition during inference. Thus...
							if not inplace:
														lowerCAmelCase__  =   ma + ma
							else:
														ma += ma

							return ma
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	List[int]						) ->							torch.Tensor:
							lowerCAmelCase__  =   -1 * len(UpperCAmelCase_						)
							lowerCAmelCase__  =   list(range(len(tensor.shape[:zero_index]						)						)						)
							return tensor.permute(first_inds + [zero_index + i for i in inds]						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int						) ->							torch.Tensor:
							return t.reshape(t.shape[:-no_dims] + (-1,)						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int					, UpperCAmelCase_      :	float = 1E-4						) ->							torch.Tensor:
							lowerCAmelCase__  =   mask.expand(*value.shape						)
							return torch.sum(mask * value					, dim=UpperCAmelCase_						) / (eps + torch.sum(UpperCAmelCase_					, dim=UpperCAmelCase_						))
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.types.Number = 2.31_25					, UpperCAmelCase_      :	torch.types.Number = 21.68_75					, UpperCAmelCase_      :	int = 64						) ->							torch.Tensor:
							lowerCAmelCase__  =   torch.linspace(UpperCAmelCase_					, UpperCAmelCase_					, no_bins - 1					, device=pts.device						)
							lowerCAmelCase__  =   torch.sqrt(torch.sum((pts.unsqueeze(-2						) - pts.unsqueeze(-3						)) ** 2					, dim=-1						)						)
							return torch.bucketize(UpperCAmelCase_					, UpperCAmelCase_						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[list], Any]					, UpperCAmelCase_      :	List[dict]						) ->							dict:
							lowerCAmelCase__  =   dicts[0]
							lowerCAmelCase__  =   {}
							for k, v in first.items():
														lowerCAmelCase__  =   [d[k] for d in dicts]
														if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
																					lowerCAmelCase__  =   dict_multimap(UpperCAmelCase_					, UpperCAmelCase_						)
														else:
																					lowerCAmelCase__  =   fn(UpperCAmelCase_						)

							return new_dict
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor						) ->							torch.Tensor:
							lowerCAmelCase__  =   v_bins.view(((1,) * len(x.shape						)) + (len(UpperCAmelCase_						),)						)
							lowerCAmelCase__  =   x[..., None] - reshaped_bins
							lowerCAmelCase__  =   torch.argmin(torch.abs(UpperCAmelCase_						)					, dim=-1						)
							return nn.functional.one_hot(UpperCAmelCase_					, num_classes=len(UpperCAmelCase_						)						).float()
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int = 0					, UpperCAmelCase_      :	int = 0						) ->							torch.Tensor:
							lowerCAmelCase__  =   []
							for i, s in enumerate(data.shape[:no_batch_dims]						):
														lowerCAmelCase__  =   torch.arange(UpperCAmelCase_						)
														lowerCAmelCase__  =   r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape						) - i - 1)))						)
														ranges.append(UpperCAmelCase_						)

							lowerCAmelCase__  =   [slice(UpperCAmelCase_						) for _ in range(len(data.shape						) - no_batch_dims						)]
							lowerCAmelCase__  =   inds
							ranges.extend(UpperCAmelCase_						)
							# Matt note: Editing this to get around the behaviour of using a list as an array index changing
							# in recent Numpy versions
							return data[tuple(UpperCAmelCase_						)]


_UpperCamelCase  						=			TypeVar("""T""")
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	Dict[Any, Union[dict, list, tuple, T]]					, UpperCAmelCase_      :	Type[T]						) ->							Dict[Any, Union[dict, list, tuple, Any]]:
							lowerCAmelCase__  =   {}
							for k, v in dic.items():
														if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
																					lowerCAmelCase__  =   dict_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)
														else:
																					lowerCAmelCase__  =   tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)

							return new_dict
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	T					, UpperCAmelCase_      :	Type[T]						) ->							Any:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	dict					, UpperCAmelCase_      :	Type[T]						) ->							dict:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	list					, UpperCAmelCase_      :	Type[T]						) ->							list:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	tuple					, UpperCAmelCase_      :	Type[T]						) ->							tuple:
							...
def 			_lowerCAmelCase(      UpperCAmelCase_      :	str					, UpperCAmelCase_      :	Dict					, UpperCAmelCase_      :	Union[str, Any]						) ->							List[str]:
							if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return dict_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return [tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						) for x in tree]
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return tuple(tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						) for x in tree						)
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return fn(UpperCAmelCase_						)
							else:
														print(type(UpperCAmelCase_						)						)
														raise ValueError("""Not supported"""						)






_UpperCamelCase  						=			partial(tree_map, leaf_type=torch.Tensor)





