
import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def 			__UpperCAmelCase		(		a_			:		Tuple    ,  a_			:		Union[str, Any]    ,  a_			:		str    ,  a_			:		Union[str, Any]    ,  a_			:		Optional[int] = 40_96				)     ->       Any:
   _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					key.shape[-3:]
   _UpperCAmelCase   	=					value.shape[-1]
   _UpperCAmelCase   	=					min(a_    ,  a_				)
   _UpperCAmelCase   	=					query / jnp.sqrt(a_				)

   @functools.partial(jax.checkpoint    ,  prevent_cse=a_				)
   def summarize_chunk(a_			:		Union[str, Any]    ,  a_			:		str    ,  a_			:		Any				):
      _UpperCAmelCase   	=					jnp.einsum("...qhd,...khd->...qhk"    ,  a_    ,  a_    ,  precision=a_				)

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

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

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

   def chunk_scanner(a_			:		Optional[Any]				):
      # julienne key array
      _UpperCAmelCase   	=					jax.lax.dynamic_slice(
          operand=a_    ,  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=a_    ,  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(a_    ,  a_    ,  a_				)

   _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					jax.lax.map(f=a_    ,  xs=jnp.arange(0    ,  a_    ,  a_				)				)

   _UpperCAmelCase   	=					jnp.max(a_    ,  axis=0    ,  keepdims=a_				)
   _UpperCAmelCase   	=					jnp.exp(chunk_max - global_max				)

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

   _UpperCAmelCase   	=					chunk_values.sum(axis=0				)
   _UpperCAmelCase   	=					jnp.expand_dims(a_    ,  -1				).sum(axis=0				)

   return all_values / all_weights
def 			__UpperCAmelCase		(		a_			:		List[str]    ,  a_			:		Union[str, Any]    ,  a_			:		List[str]    ,  a_			:		Dict=jax.lax.Precision.HIGHEST    ,  a_			:		List[Any] = 10_24    ,  a_			:		Optional[Any] = 40_96				)     ->       Dict:
   _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					query.shape[-3:]

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

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

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

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






class     __a			(					nn.Module			):
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int       =			8
       _lowerCamelCase							:							int       =			64
       _lowerCamelCase							:							float       =			0.0
       _lowerCamelCase							:							bool       =			False
       _lowerCamelCase							:							jnp.dtype       =			jnp.floataa







       def    _snake_case			(					self		):



          '''simple docstring'''

          _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_		):



          '''simple docstring'''

          _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_		):



          '''simple docstring'''

          _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		):



          '''simple docstring'''

          _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=4096 * 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     __a			(					nn.Module			):
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int
       _lowerCamelCase							:							float       =			0.0
       _lowerCamelCase							:							bool       =			False
       _lowerCamelCase							:							jnp.dtype       =			jnp.floataa
       _lowerCamelCase							:							bool       =			False







       def    _snake_case			(					self		):



          '''simple docstring'''

          _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		):



          '''simple docstring'''

          _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     __a			(					nn.Module			):
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int
       _lowerCamelCase							:							int       =			1
       _lowerCamelCase							:							float       =			0.0
       _lowerCamelCase							:							bool       =			False
       _lowerCamelCase							:							bool       =			False
       _lowerCamelCase							:							jnp.dtype       =			jnp.floataa
       _lowerCamelCase							:							bool       =			False







       def    _snake_case			(					self		):



          '''simple docstring'''

          _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		):



          '''simple docstring'''

          _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     __a			(					nn.Module			):
       _lowerCamelCase							:							int
       _lowerCamelCase							:							float       =			0.0
       _lowerCamelCase							:							jnp.dtype       =			jnp.floataa







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					FlaxGEGLU(self.dim							,self.dropout							,self.dtype		)
          _UpperCAmelCase   	=					nn.Dense(self.dim							,dtype=self.dtype		)







       def __call__(					self							,lowerCamelCase_							,lowerCamelCase_=True		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.net_a(UpperCamelCase__							,deterministic=UpperCamelCase__		)
          _UpperCAmelCase   	=					self.net_a(UpperCamelCase__		)
          return hidden_states






class     __a			(					nn.Module			):
       _lowerCamelCase							:							int
       _lowerCamelCase							:							float       =			0.0
       _lowerCamelCase							:							jnp.dtype       =			jnp.floataa







       def    _snake_case			(					self		):



          '''simple docstring'''

          _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		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.proj(UpperCamelCase__		)
          _UpperCAmelCase	,							_UpperCAmelCase   	=					jnp.split(UpperCamelCase__							,2							,axis=2		)
          return self.dropout_layer(hidden_linear * nn.gelu(UpperCamelCase__		)							,deterministic=UpperCamelCase__		)



"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
  import tensorflow as tf

  from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
  from transformers.models.albert.modeling_tf_albert import (
      TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFAlbertForMaskedLM,
      TFAlbertForMultipleChoice,
      TFAlbertForPreTraining,
      TFAlbertForQuestionAnswering,
      TFAlbertForSequenceClassification,
      TFAlbertForTokenClassification,
      TFAlbertModel,
  )






class     __a			:







       def __init__(					self							,lowerCamelCase_							,lowerCamelCase_=13							,lowerCamelCase_=7							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=99							,lowerCamelCase_=16							,lowerCamelCase_=32							,lowerCamelCase_=2							,lowerCamelCase_=4							,lowerCamelCase_=37							,lowerCamelCase_="gelu"							,lowerCamelCase_=0.1							,lowerCamelCase_=0.1							,lowerCamelCase_=512							,lowerCamelCase_=16							,lowerCamelCase_=2							,lowerCamelCase_=0.0_2							,lowerCamelCase_=3							,lowerCamelCase_=4							,lowerCamelCase_=None							,):



          '''simple docstring'''

          _UpperCAmelCase   	=					parent
          _UpperCAmelCase   	=					13
          _UpperCAmelCase   	=					7
          _UpperCAmelCase   	=					True
          _UpperCAmelCase   	=					True
          _UpperCAmelCase   	=					True
          _UpperCAmelCase   	=					True
          _UpperCAmelCase   	=					99
          _UpperCAmelCase   	=					16
          _UpperCAmelCase   	=					32
          _UpperCAmelCase   	=					2
          _UpperCAmelCase   	=					4
          _UpperCAmelCase   	=					37
          _UpperCAmelCase   	=					"gelu"
          _UpperCAmelCase   	=					0.1
          _UpperCAmelCase   	=					0.1
          _UpperCAmelCase   	=					512
          _UpperCAmelCase   	=					16
          _UpperCAmelCase   	=					2
          _UpperCAmelCase   	=					0.0_2
          _UpperCAmelCase   	=					3
          _UpperCAmelCase   	=					4
          _UpperCAmelCase   	=					None







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size		)

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

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

          _UpperCAmelCase   	=					None
          _UpperCAmelCase   	=					None
          _UpperCAmelCase   	=					None
          if self.use_labels:
             _UpperCAmelCase   	=					ids_tensor([self.batch_size]							,self.type_sequence_label_size		)
             _UpperCAmelCase   	=					ids_tensor([self.batch_size, self.seq_length]							,self.num_labels		)
             _UpperCAmelCase   	=					ids_tensor([self.batch_size]							,self.num_choices		)

          _UpperCAmelCase   	=					AlbertConfig(
              vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,embedding_size=self.embedding_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							,initializer_range=self.initializer_range							,)

          return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					TFAlbertModel(config=lowerCamelCase_		)
          # inputs = {'input_ids': input_ids,
          #           'attention_mask': input_mask,
          #           'token_type_ids': token_type_ids}
          # sequence_output, pooled_output = model(**inputs)
          _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _UpperCAmelCase   	=					model(lowerCamelCase_		)

          _UpperCAmelCase   	=					[input_ids, input_mask]
          _UpperCAmelCase   	=					model(lowerCamelCase_		)

          _UpperCAmelCase   	=					model(lowerCamelCase_		)

          self.parent.assertEqual(result.last_hidden_state.shape							,(self.batch_size, self.seq_length, self.hidden_size)		)
          self.parent.assertEqual(result.pooler_output.shape							,(self.batch_size, self.hidden_size)		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.num_labels
          _UpperCAmelCase   	=					TFAlbertForPreTraining(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertEqual(result.prediction_logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)		)
          self.parent.assertEqual(result.sop_logits.shape							,(self.batch_size, self.num_labels)		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					TFAlbertForMaskedLM(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.num_labels
          _UpperCAmelCase   	=					TFAlbertForSequenceClassification(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					TFAlbertForQuestionAnswering(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertEqual(result.start_logits.shape							,(self.batch_size, self.seq_length)		)
          self.parent.assertEqual(result.end_logits.shape							,(self.batch_size, self.seq_length)		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.num_choices
          _UpperCAmelCase   	=					TFAlbertForMultipleChoice(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
          _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
          _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
          _UpperCAmelCase   	=					{
              "input_ids": multiple_choice_inputs_ids,
              "attention_mask": multiple_choice_input_mask,
              "token_type_ids": multiple_choice_token_type_ids,
          }
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertListEqual(list(result["logits"].shape		)							,[self.batch_size, self.num_choices]		)







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.num_labels
          _UpperCAmelCase   	=					TFAlbertForTokenClassification(config=lowerCamelCase_		)
          _UpperCAmelCase   	=					{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "token_type_ids": token_type_ids,
          }
          _UpperCAmelCase   	=					model(lowerCamelCase_		)
          self.parent.assertListEqual(list(result["logits"].shape		)							,[self.batch_size, self.seq_length, self.num_labels]		)







       def    _snake_case			(					self		):



          '''simple docstring'''

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






@require_tf
class     __a			(					lowerCAmelCase_				,	lowerCAmelCase_				,	unittest.TestCase			):
       _lowerCamelCase							:							Any       =			(
           (
               TFAlbertModel,
               TFAlbertForPreTraining,
               TFAlbertForMaskedLM,
               TFAlbertForSequenceClassification,
               TFAlbertForQuestionAnswering,
               TFAlbertForTokenClassification,
               TFAlbertForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       _lowerCamelCase							:							int       =			(
           {
               'feature-extraction': TFAlbertModel,
               'fill-mask': TFAlbertForMaskedLM,
               'question-answering': TFAlbertForQuestionAnswering,
               'text-classification': TFAlbertForSequenceClassification,
               'token-classification': TFAlbertForTokenClassification,
               'zero-shot': TFAlbertForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       _lowerCamelCase							:							List[str]       =			False
       _lowerCamelCase							:							Optional[Any]       =			False







       def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_=False		):



          '''simple docstring'''

          _UpperCAmelCase   	=					super()._prepare_for_class(lowerCamelCase_							,lowerCamelCase_							,return_labels=lowerCamelCase_		)

          if return_labels:
             if model_class in get_values(lowerCamelCase_		):
                _UpperCAmelCase   	=					tf.zeros(self.model_tester.batch_size							,dtype=tf.intaa		)

          return inputs_dict







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					TFAlbertModelTester(self		)
          _UpperCAmelCase   	=					ConfigTester(self							,config_class=lowerCamelCase_							,hidden_size=37		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          self.config_tester.run_common_tests()







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_model(*lowerCamelCase_		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_for_pretraining(*lowerCamelCase_		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_for_masked_lm(*lowerCamelCase_		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCamelCase_		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCamelCase_		)







       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_albert_for_question_answering(*lowerCamelCase_		)







       @slow
       def    _snake_case			(					self		):



          '''simple docstring'''

          for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase   	=					TFAlbertModel.from_pretrained(lowerCamelCase_		)
             self.assertIsNotNone(lowerCamelCase_		)






@require_tf
class     __a			(					unittest.TestCase			):







       @slow
       def    _snake_case			(					self		):



          '''simple docstring'''

          _UpperCAmelCase   	=					TFAlbertForPreTraining.from_pretrained("albert-base-v2"		)
          _UpperCAmelCase   	=					tf.constant([[0, 1, 2, 3, 4, 5]]		)
          _UpperCAmelCase   	=					model(lowerCamelCase_		)[0]

          _UpperCAmelCase   	=					[1, 6, 3_0000]
          self.assertEqual(output.shape							,lowerCamelCase_		)

          _UpperCAmelCase   	=					tf.constant(
              [
                  [
                      [4.5_9_5_6_6_8, 0.7_4_4_6_2_7_5_4, -1.8_1_8_1_4_7],
                      [4.5_9_5_4_3_4_7, 0.7_4_5_4_1_8_4, -1.8_1_8_8_2_5_8],
                      [4.5_9_5_4_9_0_5, 0.7_4_4_8_2_3_5, -1.8_1_8_2_3_1_6],
                  ]
              ]		)
          tf.debugging.assert_near(output[:, :3, :3]							,lowerCamelCase_							,atol=1e-4		)

