






import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
				import jax.numpy as jnp

				from transformers.models.albert.modeling_flax_albert import (
				    FlaxAlbertForMaskedLM,
				    FlaxAlbertForMultipleChoice,
				    FlaxAlbertForPreTraining,
				    FlaxAlbertForQuestionAnswering,
				    FlaxAlbertForSequenceClassification,
				    FlaxAlbertForTokenClassification,
				    FlaxAlbertModel,
				)





class 						SCREAMING_SNAKE_CASE				( unittest.TestCase					):





							"""simple docstring"""

							def __init__(  self:   Optional[int]		,     __A:   Optional[Any]		,     __A:   str=13		,     __A:   str=7		,     __A:   int=True		,     __A:   Optional[Any]=True		,     __A:   int=True		,     __A:   Any=True		,     __A:   Optional[int]=99		,     __A:   Union[str, Any]=32		,     __A:   int=5		,     __A:   List[Any]=4		,     __A:   Any=37		,     __A:   List[str]="gelu"		,     __A:   Union[str, Any]=0.1		,     __A:   List[Any]=0.1		,     __A:   int=5_12		,     __A:   Optional[Any]=16		,     __A:   Any=2		,     __A:   List[str]=0.02		,     __A:   str=4		,     )    ->				Optional[Any]:
												_A								= parent
												_A								= batch_size
												_A								= seq_length
												_A								= is_training
												_A								= use_attention_mask
												_A								= use_token_type_ids
												_A								= use_labels
												_A								= vocab_size
												_A								= hidden_size
												_A								= num_hidden_layers
												_A								= num_attention_heads
												_A								= intermediate_size
												_A								= hidden_act
												_A								= hidden_dropout_prob
												_A								= attention_probs_dropout_prob
												_A								= max_position_embeddings
												_A								= type_vocab_size
												_A								= type_sequence_label_size
												_A								= initializer_range
												_A								= num_choices

							def       __A (  self:   Any   )    ->				Optional[Any]:
												_A								= ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size   )

												_A								= None
												if self.use_attention_mask:
																	_A								= random_attention_mask([self.batch_size, self.seq_length]   )

												_A								= None
												if self.use_token_type_ids:
																	_A								= ids_tensor([self.batch_size, self.seq_length]		,     self.type_vocab_size   )

												_A								= AlbertConfig(
												    vocab_size=self.vocab_size		,     hidden_size=self.hidden_size		,     num_hidden_layers=self.num_hidden_layers		,     num_attention_heads=self.num_attention_heads		,     intermediate_size=self.intermediate_size		,     hidden_act=self.hidden_act		,     hidden_dropout_prob=self.hidden_dropout_prob		,     attention_probs_dropout_prob=self.attention_probs_dropout_prob		,     max_position_embeddings=self.max_position_embeddings		,     type_vocab_size=self.type_vocab_size		,     is_decoder=_A		,     initializer_range=self.initializer_range		,     )

												return config, input_ids, token_type_ids, attention_mask



							def       __A (  self:   List[Any]   )    ->				Optional[int]:
												_A								= self.prepare_config_and_inputs()
												_A								= config_and_inputs
												_A								= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
												return config, inputs_dict





@require_flax
class 						SCREAMING_SNAKE_CASE				( a__  ,	unittest.TestCase					):





							"""simple docstring"""

							A_						  =					(
							    (
							        FlaxAlbertModel,
							        FlaxAlbertForPreTraining,
							        FlaxAlbertForMaskedLM,
							        FlaxAlbertForMultipleChoice,
							        FlaxAlbertForQuestionAnswering,
							        FlaxAlbertForSequenceClassification,
							        FlaxAlbertForTokenClassification,
							        FlaxAlbertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)

							def       __A (  self:   Any   )    ->				str:
												_A								= FlaxAlbertModelTester(self   )



							@slow
							def       __A (  self:   List[Any]   )    ->				Union[str, Any]:

												for model_class_name in self.all_model_classes:
																	_A								= model_class_name.from_pretrained('''albert-base-v2'''   )
																	_A								= model(np.ones((1, 1)   )   )
																	self.assertIsNotNone(_A   )





@require_flax
class 						SCREAMING_SNAKE_CASE				( unittest.TestCase					):





							"""simple docstring"""

							@slow
							def       __A (  self:   List[str]   )    ->				int:
												_A								= FlaxAlbertModel.from_pretrained('''albert-base-v2'''   )
												_A								= np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]   )
												_A								= np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
												_A								= model(_A		,     attention_mask=_A   )[0]
												_A								= (1, 11, 7_68)
												self.assertEqual(output.shape		,     _A   )
												_A								= np.array(
												    [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]   )

												self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]		,     _A		,     atol=1e-4   )   )














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





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   Tuple   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   Dict   )    ->				Tuple:
												_A							,_A							,_A							,_A								= hidden_states.shape
												_A								= jax.image.resize(
												    __A		,     shape=(batch, height * 2, width * 2, channels)		,     method='''nearest'''		,     )
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					jnp.floataa

							def       __A (  self:   List[str]   )    ->				Tuple:
												_A								= nn.Conv(
												    self.out_channels		,     kernel_size=(3, 3)		,     strides=(2, 2)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )



							def __call__(  self:   Union[str, Any]		,     __A:   List[Any]   )    ->				Union[str, Any]:
												# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
												# hidden_states = jnp.pad(hidden_states, pad_width=pad)
												_A								= self.conv(__A   )
												return hidden_states





class 						SCREAMING_SNAKE_CASE				( nn.Module					):





							"""simple docstring"""

							A_						  =					42
							A_						  =					None
							A_						  =					0.0
							A_						  =					None
							A_						  =					jnp.floataa

							def       __A (  self:   Dict   )    ->				Dict:
												_A								= self.in_channels if self.out_channels is None else self.out_channels

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= nn.Dense(__A		,     dtype=self.dtype   )

												_A								= nn.GroupNorm(num_groups=32		,     epsilon=1e-5   )
												_A								= nn.Dropout(self.dropout_prob   )
												_A								= nn.Conv(
												    __A		,     kernel_size=(3, 3)		,     strides=(1, 1)		,     padding=((1, 1), (1, 1))		,     dtype=self.dtype		,     )

												_A								= self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

												_A								= None
												if use_nin_shortcut:
																	_A								= nn.Conv(
																	    __A		,     kernel_size=(1, 1)		,     strides=(1, 1)		,     padding='''VALID'''		,     dtype=self.dtype		,     )



							def __call__(  self:   Dict		,     __A:   List[Any]		,     __A:   List[Any]		,     __A:   Any=True   )    ->				List[Any]:
												_A								= hidden_states
												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.conva(__A   )

												_A								= self.time_emb_proj(nn.swish(__A   )   )
												_A								= jnp.expand_dims(jnp.expand_dims(__A		,     1   )		,     1   )
												_A								= hidden_states + temb

												_A								= self.norma(__A   )
												_A								= nn.swish(__A   )
												_A								= self.dropout(__A		,     __A   )
												_A								= self.conva(__A   )

												if self.conv_shortcut is not None:
																	_A								= self.conv_shortcut(__A   )

												return hidden_states + residual



