





import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp

def      a    (   _lowercase :  int			,    _lowercase :  Optional[Any]			,    _lowercase :  int			,    _lowercase :  Dict			,    _lowercase :  int = 4_0_9_6							)  ->Any:
		UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						key.shape[-3:]
		UpperCamelCase__								=						value.shape[-1]
		UpperCamelCase__								=						min(_lowercase			,    _lowercase							)
		UpperCamelCase__								=						query / jnp.sqrt(_lowercase							)

		@functools.partial(jax.checkpoint			,    prevent_cse=_lowercase							)
		def summarize_chunk(_lowercase :  Tuple			,    _lowercase :  Dict			,    _lowercase :  Optional[int]							):
				UpperCamelCase__								=						jnp.einsum("""...qhd,...khd->...qhk"""			,    _lowercase			,    _lowercase			,    precision=_lowercase							)

				UpperCamelCase__								=						jnp.max(_lowercase			,    axis=-1			,    keepdims=_lowercase							)
				UpperCamelCase__								=						jax.lax.stop_gradient(_lowercase							)
				UpperCamelCase__								=						jnp.exp(attn_weights - max_score							)

				UpperCamelCase__								=						jnp.einsum("""...vhf,...qhv->...qhf"""			,    _lowercase			,    _lowercase			,    precision=_lowercase							)
				UpperCamelCase__								=						jnp.einsum("""...qhk->...qh"""			,    _lowercase							)

				return (exp_values, exp_weights.sum(axis=-1							), max_score)

		def chunk_scanner(_lowercase :  Tuple							):
				# julienne key array
				UpperCamelCase__								=						jax.lax.dynamic_slice(
				    operand=_lowercase			,    start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]			,    slice_sizes=list(key.shape[:-3]							) + [key_chunk_size, num_heads, k_features]			,    )

				# julienne value array
				UpperCamelCase__								=						jax.lax.dynamic_slice(
				    operand=_lowercase			,    start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]			,    slice_sizes=list(value.shape[:-3]							) + [key_chunk_size, num_heads, v_features]			,    )

				return summarize_chunk(_lowercase			,    _lowercase			,    _lowercase							)

		UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						jax.lax.map(f=_lowercase			,    xs=jnp.arange(0			,    _lowercase			,    _lowercase							)							)

		UpperCamelCase__								=						jnp.max(_lowercase			,    axis=0			,    keepdims=_lowercase							)
		UpperCamelCase__								=						jnp.exp(chunk_max - global_max							)

		chunk_values *= jnp.expand_dims(_lowercase			,    axis=-1							)
		chunk_weights *= max_diffs

		UpperCamelCase__								=						chunk_values.sum(axis=0							)
		UpperCamelCase__								=						jnp.expand_dims(_lowercase			,    -1							).sum(axis=0							)

		return all_values / all_weights

def      a    (   _lowercase :  Any			,    _lowercase :  Optional[int]			,    _lowercase :  Dict			,    _lowercase :  Dict=jax.lax.Precision.HIGHEST			,    _lowercase :  int = 1_0_2_4			,    _lowercase :  int = 4_0_9_6							)  ->Optional[int]:
		UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						query.shape[-3:]

		def chunk_scanner(_lowercase :  Optional[Any]			,    _lowercase :  Union[str, Any]							):
				# julienne query array
				UpperCamelCase__								=						jax.lax.dynamic_slice(
				    operand=_lowercase			,    start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]			,    slice_sizes=list(query.shape[:-3]							) + [min(_lowercase			,    _lowercase							), num_heads, q_features]			,    )

				return (
				    chunk_idx + query_chunk_size,  # unused ignore it
				    _query_chunk_attention(
				        query=_lowercase			,    key=_lowercase			,    value=_lowercase			,    precision=_lowercase			,    key_chunk_size=_lowercase							),
				)

		UpperCamelCase__	,      UpperCamelCase__								=						jax.lax.scan(
		    f=_lowercase			,    init=0			,    xs=_lowercase			,    length=math.ceil(num_q / query_chunk_size							)							)

		return jnp.concatenate(_lowercase			,    axis=-3							)  # fuse the chunked result back




class 	_UpperCAmelCase						( nn.Module       ):
					__A        = 42
					__A        = 8
					__A        = 64
					__A        = 0.0
					__A        = False
					__A        = jnp.floataa




					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__								=						self.dim_head * self.heads
							UpperCamelCase__								=						self.dim_head**-0.5

							# Weights were exported with old names {to_q, to_k, to_v, to_out}
							UpperCamelCase__								=						nn.Dense(_UpperCAmelCase     ,   use_bias=_UpperCAmelCase     ,   dtype=self.dtype     ,   name="""to_q"""							)
							UpperCamelCase__								=						nn.Dense(_UpperCAmelCase     ,   use_bias=_UpperCAmelCase     ,   dtype=self.dtype     ,   name="""to_k"""							)
							UpperCamelCase__								=						nn.Dense(_UpperCAmelCase     ,   use_bias=_UpperCAmelCase     ,   dtype=self.dtype     ,   name="""to_v"""							)

							UpperCamelCase__								=						nn.Dense(self.query_dim     ,   dtype=self.dtype     ,   name="""to_out_0"""							)
							UpperCamelCase__								=						nn.Dropout(rate=self.dropout							)




					def 						snake_case_		(    self     ,   __lowerCAmelCase							)    ->	int:
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						tensor.shape
							UpperCamelCase__								=						self.heads
							UpperCamelCase__								=						tensor.reshape(_UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase     ,   dim // head_size							)
							UpperCamelCase__								=						jnp.transpose(_UpperCAmelCase     ,   (0, 2, 1, 3)							)
							UpperCamelCase__								=						tensor.reshape(batch_size * head_size     ,   _UpperCAmelCase     ,   dim // head_size							)
							return tensor




					def 						snake_case_		(    self     ,   __lowerCAmelCase							)    ->	int:
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						tensor.shape
							UpperCamelCase__								=						self.heads
							UpperCamelCase__								=						tensor.reshape(batch_size // head_size     ,   _UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase							)
							UpperCamelCase__								=						jnp.transpose(_UpperCAmelCase     ,   (0, 2, 1, 3)							)
							UpperCamelCase__								=						tensor.reshape(batch_size // head_size     ,   _UpperCAmelCase     ,   dim * head_size							)
							return tensor




					def __call__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase=None     ,   __lowerCAmelCase=True							)    ->	int:
							UpperCamelCase__								=						hidden_states if context is None else context

							UpperCamelCase__								=						self.query(_UpperCAmelCase							)
							UpperCamelCase__								=						self.key(_UpperCAmelCase							)
							UpperCamelCase__								=						self.value(_UpperCAmelCase							)

							UpperCamelCase__								=						self.reshape_heads_to_batch_dim(_UpperCAmelCase							)
							UpperCamelCase__								=						self.reshape_heads_to_batch_dim(_UpperCAmelCase							)
							UpperCamelCase__								=						self.reshape_heads_to_batch_dim(_UpperCAmelCase							)

							if self.use_memory_efficient_attention:
									UpperCamelCase__								=						query_states.transpose(1     ,   0     ,   2							)
									UpperCamelCase__								=						key_states.transpose(1     ,   0     ,   2							)
									UpperCamelCase__								=						value_states.transpose(1     ,   0     ,   2							)

									# this if statement create a chunk size for each layer of the unet
									# the chunk size is equal to the query_length dimension of the deepest layer of the unet

									UpperCamelCase__								=						query_states.shape[-3]
									if flatten_latent_dim % 64 == 0:
											UpperCamelCase__								=						int(flatten_latent_dim / 64							)
									elif flatten_latent_dim % 16 == 0:
											UpperCamelCase__								=						int(flatten_latent_dim / 16							)
									elif flatten_latent_dim % 4 == 0:
											UpperCamelCase__								=						int(flatten_latent_dim / 4							)
									else:
											UpperCamelCase__								=						int(_UpperCAmelCase							)

									UpperCamelCase__								=						jax_memory_efficient_attention(
									    _UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase     ,   query_chunk_size=_UpperCAmelCase     ,   key_chunk_size=4_096 * 4							)

									UpperCamelCase__								=						hidden_states.transpose(1     ,   0     ,   2							)
							else:
									# compute attentions
									UpperCamelCase__								=						jnp.einsum("""b i d, b j d->b i j"""     ,   _UpperCAmelCase     ,   _UpperCAmelCase							)
									UpperCamelCase__								=						attention_scores * self.scale
									UpperCamelCase__								=						nn.softmax(_UpperCAmelCase     ,   axis=2							)

									# attend to values
									UpperCamelCase__								=						jnp.einsum("""b i j, b j d -> b i d"""     ,   _UpperCAmelCase     ,   _UpperCAmelCase							)

							UpperCamelCase__								=						self.reshape_batch_dim_to_heads(_UpperCAmelCase							)
							UpperCamelCase__								=						self.proj_attn(_UpperCAmelCase							)
							return self.dropout_layer(_UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)




class 	_UpperCAmelCase						( nn.Module       ):
					__A        = 42
					__A        = 42
					__A        = 42
					__A        = 0.0
					__A        = False
					__A        = jnp.floataa
					__A        = False




					def 						snake_case_		(    self							)    ->	int:
							# self attention (or cross_attention if only_cross_attention is True)
							UpperCamelCase__								=						FlaxAttention(
							    self.dim     ,   self.n_heads     ,   self.d_head     ,   self.dropout     ,   self.use_memory_efficient_attention     ,   dtype=self.dtype							)
							# cross attention
							UpperCamelCase__								=						FlaxAttention(
							    self.dim     ,   self.n_heads     ,   self.d_head     ,   self.dropout     ,   self.use_memory_efficient_attention     ,   dtype=self.dtype							)
							UpperCamelCase__								=						FlaxFeedForward(dim=self.dim     ,   dropout=self.dropout     ,   dtype=self.dtype							)
							UpperCamelCase__								=						nn.LayerNorm(epsilon=1E-5     ,   dtype=self.dtype							)
							UpperCamelCase__								=						nn.LayerNorm(epsilon=1E-5     ,   dtype=self.dtype							)
							UpperCamelCase__								=						nn.LayerNorm(epsilon=1E-5     ,   dtype=self.dtype							)
							UpperCamelCase__								=						nn.Dropout(rate=self.dropout							)




					def __call__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase=True							)    ->	List[Any]:
							# self attention
							UpperCamelCase__								=						hidden_states
							if self.only_cross_attention:
									UpperCamelCase__								=						self.attna(self.norma(_UpperCAmelCase							)     ,   _UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)
							else:
									UpperCamelCase__								=						self.attna(self.norma(_UpperCAmelCase							)     ,   deterministic=_UpperCAmelCase							)
							UpperCamelCase__								=						hidden_states + residual

							# cross attention
							UpperCamelCase__								=						hidden_states
							UpperCamelCase__								=						self.attna(self.norma(_UpperCAmelCase							)     ,   _UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)
							UpperCamelCase__								=						hidden_states + residual

							# feed forward
							UpperCamelCase__								=						hidden_states
							UpperCamelCase__								=						self.ff(self.norma(_UpperCAmelCase							)     ,   deterministic=_UpperCAmelCase							)
							UpperCamelCase__								=						hidden_states + residual

							return self.dropout_layer(_UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)




class 	_UpperCAmelCase						( nn.Module       ):
					__A        = 42
					__A        = 42
					__A        = 42
					__A        = 1
					__A        = 0.0
					__A        = False
					__A        = False
					__A        = jnp.floataa
					__A        = False




					def 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						nn.GroupNorm(num_groups=32     ,   epsilon=1E-5							)

							UpperCamelCase__								=						self.n_heads * self.d_head
							if self.use_linear_projection:
									UpperCamelCase__								=						nn.Dense(_UpperCAmelCase     ,   dtype=self.dtype							)
							else:
									UpperCamelCase__								=						nn.Conv(
									    _UpperCAmelCase     ,   kernel_size=(1, 1)     ,   strides=(1, 1)     ,   padding="""VALID"""     ,   dtype=self.dtype     ,   )

							UpperCamelCase__								=						[
							    FlaxBasicTransformerBlock(
							        _UpperCAmelCase     ,   self.n_heads     ,   self.d_head     ,   dropout=self.dropout     ,   only_cross_attention=self.only_cross_attention     ,   dtype=self.dtype     ,   use_memory_efficient_attention=self.use_memory_efficient_attention     ,   )
							    for _ in range(self.depth							)
							]

							if self.use_linear_projection:
									UpperCamelCase__								=						nn.Dense(_UpperCAmelCase     ,   dtype=self.dtype							)
							else:
									UpperCamelCase__								=						nn.Conv(
									    _UpperCAmelCase     ,   kernel_size=(1, 1)     ,   strides=(1, 1)     ,   padding="""VALID"""     ,   dtype=self.dtype     ,   )

							UpperCamelCase__								=						nn.Dropout(rate=self.dropout							)




					def __call__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase     ,   __lowerCAmelCase=True							)    ->	Dict:
							UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__								=						hidden_states.shape
							UpperCamelCase__								=						hidden_states
							UpperCamelCase__								=						self.norm(_UpperCAmelCase							)
							if self.use_linear_projection:
									UpperCamelCase__								=						hidden_states.reshape(_UpperCAmelCase     ,   height * width     ,   _UpperCAmelCase							)
									UpperCamelCase__								=						self.proj_in(_UpperCAmelCase							)
							else:
									UpperCamelCase__								=						self.proj_in(_UpperCAmelCase							)
									UpperCamelCase__								=						hidden_states.reshape(_UpperCAmelCase     ,   height * width     ,   _UpperCAmelCase							)

							for transformer_block in self.transformer_blocks:
									UpperCamelCase__								=						transformer_block(_UpperCAmelCase     ,   _UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)

							if self.use_linear_projection:
									UpperCamelCase__								=						self.proj_out(_UpperCAmelCase							)
									UpperCamelCase__								=						hidden_states.reshape(_UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase							)
							else:
									UpperCamelCase__								=						hidden_states.reshape(_UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase     ,   _UpperCAmelCase							)
									UpperCamelCase__								=						self.proj_out(_UpperCAmelCase							)

							UpperCamelCase__								=						hidden_states + residual
							return self.dropout_layer(_UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)




class 	_UpperCAmelCase						( nn.Module       ):
					__A        = 42
					__A        = 0.0
					__A        = jnp.floataa




					def 						snake_case_		(    self							)    ->	str:
							# The second linear layer needs to be called
							# net_2 for now to match the index of the Sequential layer
							UpperCamelCase__								=						FlaxGEGLU(self.dim     ,   self.dropout     ,   self.dtype							)
							UpperCamelCase__								=						nn.Dense(self.dim     ,   dtype=self.dtype							)




					def __call__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase=True							)    ->	Optional[Any]:
							UpperCamelCase__								=						self.net_a(_UpperCAmelCase     ,   deterministic=_UpperCAmelCase							)
							UpperCamelCase__								=						self.net_a(_UpperCAmelCase							)
							return hidden_states




class 	_UpperCAmelCase						( nn.Module       ):
					__A        = 42
					__A        = 0.0
					__A        = jnp.floataa




					def 						snake_case_		(    self							)    ->	str:
							UpperCamelCase__								=						self.dim * 4
							UpperCamelCase__								=						nn.Dense(inner_dim * 2     ,   dtype=self.dtype							)
							UpperCamelCase__								=						nn.Dropout(rate=self.dropout							)




					def __call__(    self     ,   __lowerCAmelCase     ,   __lowerCAmelCase=True							)    ->	Dict:
							UpperCamelCase__								=						self.proj(_UpperCAmelCase							)
							UpperCamelCase__	,      UpperCamelCase__								=						jnp.split(_UpperCAmelCase     ,   2     ,   axis=2							)
							return self.dropout_layer(hidden_linear * nn.gelu(_UpperCAmelCase							)     ,   deterministic=_UpperCAmelCase							)













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							, _UpperCAmelCase							, _UpperCAmelCase=14							, _UpperCAmelCase=7							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=True							, _UpperCAmelCase=99							, _UpperCAmelCase=32							, _UpperCAmelCase=5							, _UpperCAmelCase=4							, _UpperCAmelCase=37							, _UpperCAmelCase="gelu"							, _UpperCAmelCase=0.1							, _UpperCAmelCase=0.1							, _UpperCAmelCase=512							, _UpperCAmelCase=0.02							, _UpperCAmelCase=None							, )		->				str:
        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 						lowerCamelCase							(	self  )		->				List[str]:
        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=_UpperCAmelCase							, 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  )		->				Tuple:
        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 						lowerCamelCase							(	self  )		->				int:
        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 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				Optional[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        lowerCamelCase									=model.init_cache(input_ids.shape[0]							, _UpperCAmelCase  )
        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=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

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

        lowerCamelCase									=model(_UpperCAmelCase  )

        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 						lowerCamelCase							(	self							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )		->				List[Any]:
        lowerCamelCase									=20
        lowerCamelCase									=model_class_name(_UpperCAmelCase  )

        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]							, _UpperCAmelCase  )
        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=_UpperCAmelCase							, past_key_values=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )
        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=_UpperCAmelCase							, position_ids=_UpperCAmelCase							, )

        lowerCamelCase									=model(_UpperCAmelCase							, attention_mask=_UpperCAmelCase  )
        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_				(  UpperCamelCase				,     UpperCamelCase				,     unittest.TestCase ):
     __lowerCamelCase       :    Tuple 	=     (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
     __lowerCamelCase       :    Dict 	=     (FlaxXGLMForCausalLM,) if is_flax_available() else ()





     def 						lowerCamelCase							(	self  )		->				List[str]:
        lowerCamelCase									=FlaxXGLMModelTester(self  )





     def 						lowerCamelCase							(	self  )		->				List[str]:
        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(_UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     def 						lowerCamelCase							(	self  )		->				Tuple:
        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(
               _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase							, _UpperCAmelCase  )





     @slow
     def 						lowerCamelCase							(	self  )		->				Tuple:
        lowerCamelCase									=XGLMTokenizer.from_pretrained("""XGLM"""							, padding_side="""left"""  )
        lowerCamelCase									=tokenizer(["""Hello this is a long string""", """Hey"""]							, return_tensors="""np"""							, padding=_UpperCAmelCase							, truncation=_UpperCAmelCase  )

        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(_UpperCAmelCase							, skip_special_tokens=_UpperCAmelCase  )

        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(_UpperCAmelCase							, _UpperCAmelCase  )





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Dict:
        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(_UpperCAmelCase							, _UpperCAmelCase  )
              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(_UpperCAmelCase							, _UpperCAmelCase  )

              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(_UpperCAmelCase  ):
                 lowerCamelCase									=0
                 lowerCamelCase									=1
                 lowerCamelCase									=0
                 lowerCamelCase									=1
              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).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(_UpperCAmelCase							, dtype=jnp.floataa  )
              lowerCamelCase									=convert_pytorch_state_dict_to_flax(pt_model.state_dict()							, _UpperCAmelCase  )
              lowerCamelCase									=fx_state

              with torch.no_grad():
                 lowerCamelCase									=pt_model(**_UpperCAmelCase  ).to_tuple()

              lowerCamelCase									=fx_model(**_UpperCAmelCase  ).to_tuple()
              self.assertEqual(len(_UpperCAmelCase  )							, len(_UpperCAmelCase  )							, """Output lengths differ between Flax and PyTorch"""  )
              for fx_output, pt_output in zip(_UpperCAmelCase							, _UpperCAmelCase  ):
                 self.assert_almost_equals(fx_output[:, -1]							, pt_output[:, -1].numpy()							, 4E-2  )

              with tempfile.TemporaryDirectory() as tmpdirname:
                 pt_model.save_pretrained(_UpperCAmelCase  )
                 lowerCamelCase									=model_class.from_pretrained(_UpperCAmelCase							, from_pt=_UpperCAmelCase  )

              lowerCamelCase									=fx_model_loaded(**_UpperCAmelCase  ).to_tuple()
              self.assertEqual(
                  len(_UpperCAmelCase  )							, len(_UpperCAmelCase  )							, """Output lengths differ between Flax and PyTorch"""  )
              for fx_output_loaded, pt_output in zip(_UpperCAmelCase							, _UpperCAmelCase  ):
                 self.assert_almost_equals(fx_output_loaded[:, -1]							, pt_output[:, -1].numpy()							, 4E-2  )





     @is_pt_flax_cross_test
     def 						lowerCamelCase							(	self  )		->				Optional[Any]:
        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(_UpperCAmelCase							, _UpperCAmelCase  )
              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(_UpperCAmelCase							, _UpperCAmelCase  )

              lowerCamelCase									=pt_model_class(_UpperCAmelCase  ).eval()
              lowerCamelCase									=False
              lowerCamelCase									=model_class(_UpperCAmelCase							, dtype=jnp.floataa  )

              lowerCamelCase									=load_flax_weights_in_pytorch_model(_UpperCAmelCase							, 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(_UpperCAmelCase  ):
                 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(**_UpperCAmelCase  ).to_tuple()

              lowerCamelCase									=fx_model(**_UpperCAmelCase  ).to_tuple()
              self.assertEqual(len(_UpperCAmelCase  )							, len(_UpperCAmelCase  )							, """Output lengths differ between Flax and PyTorch"""  )
              for fx_output, pt_output in zip(_UpperCAmelCase							, _UpperCAmelCase  ):
                 self.assert_almost_equals(fx_output[:, -1]							, pt_output[:, -1].numpy()							, 4E-2  )

              with tempfile.TemporaryDirectory() as tmpdirname:
                 fx_model.save_pretrained(_UpperCAmelCase  )
                 lowerCamelCase									=pt_model_class.from_pretrained(_UpperCAmelCase							, from_flax=_UpperCAmelCase  )

              with torch.no_grad():
                 lowerCamelCase									=pt_model_loaded(**_UpperCAmelCase  ).to_tuple()

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





     @slow
     def 						lowerCamelCase							(	self  )		->				int:
        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(_UpperCAmelCase  )