code: 
  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__		)
  
  
style_context: 
  """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		)
  
  