

'''simple docstring'''
from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)
@flax.struct.dataclass
class 					_snake_case		(__SCREAMING_SNAKE_CASE):
			__A  :	jnp.ndarray




@flax_register_to_config
class 					_snake_case		(nn.Module			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE):
			__A  :	int        =32
			__A  :	int        =4
			__A  :	int        =4
			__A  :	Tuple[str]        =(
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			__A  :	Tuple[str]        =("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
			__A  :	Union[bool, Tuple[bool]]        =False
			__A  :	Tuple[int]        =(3_20, 6_40, 12_80, 12_80)
			__A  :	int        =2
			__A  :	Union[int, Tuple[int]]        =8
			__A  :	Optional[Union[int, Tuple[int]]]        =None
			__A  :	int        =12_80
			__A  :	float        =0.0
			__A  :	bool        =False
			__A  :	jnp.dtype        =jnp.floataa
			__A  :	bool        =True
			__A  :	int        =0
			__A  :	bool        =False






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				# init input tensors
				UpperCAmelCase_     :      List[str]						=  (1, self.in_channels, self.sample_size, self.sample_size)
				UpperCAmelCase_     :      Tuple						=  jnp.zeros(_snake_case			,dtype=jnp.floataa     )
				UpperCAmelCase_     :      Optional[Any]						=  jnp.ones((1,)			,dtype=jnp.intaa     )
				UpperCAmelCase_     :      Optional[Any]						=  jnp.zeros((1, 1, self.cross_attention_dim)			,dtype=jnp.floataa     )

				UpperCAmelCase_					,					UpperCAmelCase_     :      Tuple						=  jax.random.split(_snake_case     )
				UpperCAmelCase_     :      int						=  {"params": params_rng, "dropout": dropout_rng}

				return self.init(_snake_case			,_snake_case			,_snake_case			,_snake_case     )["params"]






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      List[str]						=  self.block_out_channels
				UpperCAmelCase_     :      Tuple						=  block_out_channels[0] * 4

				if self.num_attention_heads is not None:
					raise ValueError(
					    "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."     )

				# If `num_attention_heads` is not defined (which is the case for most models)
				# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
				# The reason for this behavior is to correct for incorrectly named variables that were introduced
				# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
				# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
				# which is why we correct for the naming here.
				UpperCAmelCase_     :      Union[str, Any]						=  self.num_attention_heads or self.attention_head_dim

				# input
				UpperCAmelCase_     :      List[Any]						=  nn.Conv(
				    block_out_channels[0]			,kernel_size=(3, 3)			,strides=(1, 1)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)

				# time
				UpperCAmelCase_     :      Tuple						=  FlaxTimesteps(
				    block_out_channels[0]			,flip_sin_to_cos=self.flip_sin_to_cos			,freq_shift=self.config.freq_shift     )
				UpperCAmelCase_     :      int						=  FlaxTimestepEmbedding(_snake_case			,dtype=self.dtype     )

				UpperCAmelCase_     :      str						=  self.only_cross_attention
				if isinstance(_snake_case			,_snake_case     ):
					UpperCAmelCase_     :      Dict						=  (only_cross_attention,) * len(self.down_block_types     )

				if isinstance(_snake_case			,_snake_case     ):
					UpperCAmelCase_     :      List[str]						=  (num_attention_heads,) * len(self.down_block_types     )

				# down
				UpperCAmelCase_     :      int						=  []
				UpperCAmelCase_     :      Optional[Any]						=  block_out_channels[0]
				for i, down_block_type in enumerate(self.down_block_types     ):
					UpperCAmelCase_     :      Any						=  output_channel
					UpperCAmelCase_     :      List[Any]						=  block_out_channels[i]
					UpperCAmelCase_     :      str						=  i == len(_snake_case     ) - 1

					if down_block_type == "CrossAttnDownBlock2D":
						UpperCAmelCase_     :      Optional[Any]						=  FlaxCrossAttnDownBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,dropout=self.dropout			,num_layers=self.layers_per_block			,num_attention_heads=num_attention_heads[i]			,add_downsample=not is_final_block			,use_linear_projection=self.use_linear_projection			,only_cross_attention=only_cross_attention[i]			,use_memory_efficient_attention=self.use_memory_efficient_attention			,dtype=self.dtype			,)
					else:
						UpperCAmelCase_     :      Optional[int]						=  FlaxDownBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,dropout=self.dropout			,num_layers=self.layers_per_block			,add_downsample=not is_final_block			,dtype=self.dtype			,)

					down_blocks.append(_snake_case     )
				UpperCAmelCase_     :      Any						=  down_blocks

				# mid
				UpperCAmelCase_     :      Union[str, Any]						=  FlaxUNetMidBlockaDCrossAttn(
				    in_channels=block_out_channels[-1]			,dropout=self.dropout			,num_attention_heads=num_attention_heads[-1]			,use_linear_projection=self.use_linear_projection			,use_memory_efficient_attention=self.use_memory_efficient_attention			,dtype=self.dtype			,)

				# up
				UpperCAmelCase_     :      str						=  []
				UpperCAmelCase_     :      Union[str, Any]						=  list(reversed(_snake_case     )     )
				UpperCAmelCase_     :      Union[str, Any]						=  list(reversed(_snake_case     )     )
				UpperCAmelCase_     :      str						=  list(reversed(_snake_case     )     )
				UpperCAmelCase_     :      Any						=  reversed_block_out_channels[0]
				for i, up_block_type in enumerate(self.up_block_types     ):
					UpperCAmelCase_     :      Tuple						=  output_channel
					UpperCAmelCase_     :      Any						=  reversed_block_out_channels[i]
					UpperCAmelCase_     :      List[Any]						=  reversed_block_out_channels[min(i + 1			,len(_snake_case     ) - 1     )]

					UpperCAmelCase_     :      Optional[Any]						=  i == len(_snake_case     ) - 1

					if up_block_type == "CrossAttnUpBlock2D":
						UpperCAmelCase_     :      Dict						=  FlaxCrossAttnUpBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,prev_output_channel=_snake_case			,num_layers=self.layers_per_block + 1			,num_attention_heads=reversed_num_attention_heads[i]			,add_upsample=not is_final_block			,dropout=self.dropout			,use_linear_projection=self.use_linear_projection			,only_cross_attention=only_cross_attention[i]			,use_memory_efficient_attention=self.use_memory_efficient_attention			,dtype=self.dtype			,)
					else:
						UpperCAmelCase_     :      str						=  FlaxUpBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,prev_output_channel=_snake_case			,num_layers=self.layers_per_block + 1			,add_upsample=not is_final_block			,dropout=self.dropout			,dtype=self.dtype			,)

					up_blocks.append(_snake_case     )
					UpperCAmelCase_     :      List[Any]						=  output_channel
				UpperCAmelCase_     :      List[str]						=  up_blocks

				# out
				UpperCAmelCase_     :      Optional[int]						=  nn.GroupNorm(num_groups=32			,epsilon=1E-5     )
				UpperCAmelCase_     :      str						=  nn.Conv(
				    self.out_channels			,kernel_size=(3, 3)			,strides=(1, 1)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)






			def __call__(  self			,_snake_case			,_snake_case			,_snake_case			,_snake_case=None			,_snake_case=None			,_snake_case = True			,_snake_case = False			,):
				# 1. time
				if not isinstance(_snake_case			,jnp.ndarray     ):
					UpperCAmelCase_     :      Optional[int]						=  jnp.array([timesteps]			,dtype=jnp.intaa     )
				elif isinstance(_snake_case			,jnp.ndarray     ) and len(timesteps.shape     ) == 0:
					UpperCAmelCase_     :      List[str]						=  timesteps.astype(dtype=jnp.floataa     )
					UpperCAmelCase_     :      Optional[int]						=  jnp.expand_dims(_snake_case			,0     )

				UpperCAmelCase_     :      Union[str, Any]						=  self.time_proj(_snake_case     )
				UpperCAmelCase_     :      Optional[Any]						=  self.time_embedding(_snake_case     )

				# 2. pre-process
				UpperCAmelCase_     :      Optional[int]						=  jnp.transpose(_snake_case			,(0, 2, 3, 1)     )
				UpperCAmelCase_     :      Union[str, Any]						=  self.conv_in(_snake_case     )

				# 3. down
				UpperCAmelCase_     :      int						=  (sample,)
				for down_block in self.down_blocks:
					if isinstance(_snake_case			,_snake_case     ):
						UpperCAmelCase_					,					UpperCAmelCase_     :      Union[str, Any]						=  down_block(_snake_case			,_snake_case			,_snake_case			,deterministic=not train     )
					else:
						UpperCAmelCase_					,					UpperCAmelCase_     :      Optional[Any]						=  down_block(_snake_case			,_snake_case			,deterministic=not train     )
					down_block_res_samples += res_samples

				if down_block_additional_residuals is not None:
					UpperCAmelCase_     :      int						=  ()

					for down_block_res_sample, down_block_additional_residual in zip(
					    _snake_case			,_snake_case     ):
						down_block_res_sample += down_block_additional_residual
						new_down_block_res_samples += (down_block_res_sample,)

					UpperCAmelCase_     :      Optional[Any]						=  new_down_block_res_samples

				# 4. mid
				UpperCAmelCase_     :      Dict						=  self.mid_block(_snake_case			,_snake_case			,_snake_case			,deterministic=not train     )

				if mid_block_additional_residual is not None:
					sample += mid_block_additional_residual

				# 5. up
				for up_block in self.up_blocks:
					UpperCAmelCase_     :      List[str]						=  down_block_res_samples[-(self.layers_per_block + 1) :]
					UpperCAmelCase_     :      Tuple						=  down_block_res_samples[: -(self.layers_per_block + 1)]
					if isinstance(_snake_case			,_snake_case     ):
						UpperCAmelCase_     :      Dict						=  up_block(
						    _snake_case			,temb=_snake_case			,encoder_hidden_states=_snake_case			,res_hidden_states_tuple=_snake_case			,deterministic=not train			,)
					else:
						UpperCAmelCase_     :      Tuple						=  up_block(_snake_case			,temb=_snake_case			,res_hidden_states_tuple=_snake_case			,deterministic=not train     )

        # 6. post-process
				UpperCAmelCase_     :      Optional[int]						=  self.conv_norm_out(_snake_case     )
				UpperCAmelCase_     :      List[Any]						=  nn.silu(_snake_case     )
				UpperCAmelCase_     :      str						=  self.conv_out(_snake_case     )
				UpperCAmelCase_     :      str						=  jnp.transpose(_snake_case			,(0, 3, 1, 2)     )

				if not return_dict:
					return (sample,)

				return FlaxUNetaDConditionOutput(sample=_snake_case     )












'''simple docstring'''







from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
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 (
		    TFConvBertForMaskedLM,
		    TFConvBertForMultipleChoice,
		    TFConvBertForQuestionAnswering,
		    TFConvBertForSequenceClassification,
		    TFConvBertForTokenClassification,
		    TFConvBertModel,
		)





class 						SCREAMING_SNAKE_CASE       :






			def __init__(			self   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      str=13	,    __SCREAMING_SNAKE_CASE   :      Optional[int]=7	,    __SCREAMING_SNAKE_CASE   :      List[str]=True	,    __SCREAMING_SNAKE_CASE   :      Optional[int]=True	,    __SCREAMING_SNAKE_CASE   :      Any=True	,    __SCREAMING_SNAKE_CASE   :      int=True	,    __SCREAMING_SNAKE_CASE   :      List[str]=99	,    __SCREAMING_SNAKE_CASE   :      Optional[int]=32	,    __SCREAMING_SNAKE_CASE   :      Tuple=2	,    __SCREAMING_SNAKE_CASE   :      str=4	,    __SCREAMING_SNAKE_CASE   :      int=37	,    __SCREAMING_SNAKE_CASE   :      Any="gelu"	,    __SCREAMING_SNAKE_CASE   :      List[Any]=0.1	,    __SCREAMING_SNAKE_CASE   :      List[Any]=0.1	,    __SCREAMING_SNAKE_CASE   :      Tuple=512	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]=16	,    __SCREAMING_SNAKE_CASE   :      Dict=2	,    __SCREAMING_SNAKE_CASE   :      List[Any]=0.02	,    __SCREAMING_SNAKE_CASE   :      List[Any]=3	,    __SCREAMING_SNAKE_CASE   :      List[Any]=4	,    __SCREAMING_SNAKE_CASE   :      Dict=None	,    )      ->      Tuple:
							a_		:	Any									=   parent
							a_		:	List[str]									=   13
							a_		:	Dict									=   7
							a_		:	List[str]									=   True
							a_		:	str									=   True
							a_		:	List[str]									=   True
							a_		:	Any									=   True
							a_		:	List[Any]									=   99
							a_		:	Optional[Any]									=   384
							a_		:	List[str]									=   2
							a_		:	Optional[int]									=   4
							a_		:	Optional[int]									=   37
							a_		:	Union[str, Any]									=   '''gelu'''
							a_		:	List[str]									=   0.1
							a_		:	str									=   0.1
							a_		:	Optional[Any]									=   512
							a_		:	Union[str, Any]									=   16
							a_		:	Optional[Any]									=   2
							a_		:	Optional[int]									=   0.02
							a_		:	Union[str, Any]									=   3
							a_		:	Union[str, Any]									=   4
							a_		:	List[Any]									=   128
							a_		:	Optional[Any]									=   2
							a_		:	Any									=   9
							a_		:	Optional[int]									=   1
							a_		:	int									=   None






			def       SCREAMING_SNAKE_CASE   (			self   :      Dict       )      ->      Union[str, Any]:
							a_		:	Dict									=   ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size       )

							a_		:	int									=   None
							if self.use_input_mask:
											a_		:	Dict									=   random_attention_mask([self.batch_size, self.seq_length]       )

							a_		:	Optional[int]									=   None
							if self.use_token_type_ids:
											a_		:	Any									=   ids_tensor([self.batch_size, self.seq_length]	,    self.type_vocab_size       )

							a_		:	Optional[Any]									=   None
							a_		:	Optional[Any]									=   None
							a_		:	Union[str, Any]									=   None
							if self.use_labels:
											a_		:	List[str]									=   ids_tensor([self.batch_size]	,    self.type_sequence_label_size       )
											a_		:	int									=   ids_tensor([self.batch_size, self.seq_length]	,    self.num_labels       )
											a_		:	Tuple									=   ids_tensor([self.batch_size]	,    self.num_choices       )

							a_		:	Union[str, Any]									=   ConvBertConfig(
							    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	,    initializer_range=self.initializer_range	,    return_dict=__SCREAMING_SNAKE_CASE	,    )

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






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      int	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]       )      ->      Union[str, Any]:
							a_		:	Tuple									=   TFConvBertModel(config=__SCREAMING_SNAKE_CASE       )
							a_		:	str									=   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

							a_		:	Optional[Any]									=   [input_ids, input_mask]
							a_		:	Tuple									=   model(__SCREAMING_SNAKE_CASE       )

							a_		:	int									=   model(__SCREAMING_SNAKE_CASE       )

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






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Dict       )      ->      str:
							a_		:	List[Any]									=   TFConvBertForMaskedLM(config=__SCREAMING_SNAKE_CASE       )
							a_		:	Dict									=   {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}
							a_		:	str									=   model(__SCREAMING_SNAKE_CASE       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size)       )






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[int]	,    __SCREAMING_SNAKE_CASE   :      str	,    __SCREAMING_SNAKE_CASE   :      int	,    __SCREAMING_SNAKE_CASE   :      str	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Dict	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      int       )      ->      Any:
							a_		:	Union[str, Any]									=   self.num_labels
							a_		:	List[str]									=   TFConvBertForSequenceClassification(config=__SCREAMING_SNAKE_CASE       )
							a_		:	Optional[Any]									=   {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}

							a_		:	Tuple									=   model(__SCREAMING_SNAKE_CASE       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      str	,    __SCREAMING_SNAKE_CASE   :      str	,    __SCREAMING_SNAKE_CASE   :      Dict	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Tuple       )      ->      Optional[int]:
							a_		:	Optional[Any]									=   self.num_choices
							a_		:	int									=   TFConvBertForMultipleChoice(config=__SCREAMING_SNAKE_CASE       )
							a_		:	List[str]									=   tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE	,    1       )	,    (1, self.num_choices, 1)       )
							a_		:	Tuple									=   tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE	,    1       )	,    (1, self.num_choices, 1)       )
							a_		:	List[Any]									=   tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE	,    1       )	,    (1, self.num_choices, 1)       )
							a_		:	Union[str, Any]									=   {
							    '''input_ids''': multiple_choice_inputs_ids,
							    '''attention_mask''': multiple_choice_input_mask,
							    '''token_type_ids''': multiple_choice_token_type_ids,
							}
							a_		:	Optional[int]									=   model(__SCREAMING_SNAKE_CASE       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_choices)       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Dict	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      Tuple       )      ->      str:
							a_		:	Optional[Any]									=   self.num_labels
							a_		:	Optional[int]									=   TFConvBertForTokenClassification(config=__SCREAMING_SNAKE_CASE       )
							a_		:	Tuple									=   {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}
							a_		:	Tuple									=   model(__SCREAMING_SNAKE_CASE       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.num_labels)       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      str	,    __SCREAMING_SNAKE_CASE   :      Optional[int]	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]       )      ->      Tuple:
							a_		:	List[Any]									=   TFConvBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE       )
							a_		:	Optional[Any]									=   {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''token_type_ids''': token_type_ids,
							}

							a_		:	Dict									=   model(__SCREAMING_SNAKE_CASE       )
							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       SCREAMING_SNAKE_CASE   (			self   :      List[str]       )      ->      Any:
							a_		:	List[Any]									=   self.prepare_config_and_inputs()
							(
							    (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      
							)		:	List[Any]									=   config_and_inputs
							a_		:	Any									=   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
							return config, inputs_dict





@require_tf
class 						SCREAMING_SNAKE_CASE       (       SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, unittest.TestCase							):
			snake_case__	=					(
			    (
			        TFConvBertModel,
			        TFConvBertForMaskedLM,
			        TFConvBertForQuestionAnswering,
			        TFConvBertForSequenceClassification,
			        TFConvBertForTokenClassification,
			        TFConvBertForMultipleChoice,
			    )
			    if is_tf_available()
			    else ()
			)
			snake_case__	=					(
			    {
			        "feature-extraction": TFConvBertModel,
			        "fill-mask": TFConvBertForMaskedLM,
			        "question-answering": TFConvBertForQuestionAnswering,
			        "text-classification": TFConvBertForSequenceClassification,
			        "token-classification": TFConvBertForTokenClassification,
			        "zero-shot": TFConvBertForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			snake_case__	=					False
			snake_case__	=					False
			snake_case__	=					False






			def       SCREAMING_SNAKE_CASE   (			self   :      List[str]       )      ->      List[Any]:
							a_		:	Optional[int]									=   TFConvBertModelTester(self       )
							a_		:	Any									=   ConfigTester(self	,    config_class=__SCREAMING_SNAKE_CASE	,    hidden_size=37       )






			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      Optional[Any]:
							self.config_tester.run_common_tests()






			def       SCREAMING_SNAKE_CASE   (			self   :      Tuple       )      ->      str:
							a_		:	Any									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[Any]       )      ->      List[str]:
							a_		:	str									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      Any       )      ->      str:
							a_		:	Dict									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      Optional[Any]:
							a_		:	Optional[Any]									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      Any       )      ->      str:
							a_		:	List[str]									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      Optional[int]:
							a_		:	Any									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE       )






			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[int]       )      ->      Union[str, Any]:
							a_					,      a_		:	Optional[Any]									=   self.model_tester.prepare_config_and_inputs_for_common()
							a_		:	Any									=   True
							a_		:	Optional[int]									=   True

							if hasattr(__SCREAMING_SNAKE_CASE	,    '''use_cache'''       ):
											a_		:	Dict									=   True

							a_		:	Any									=   getattr(self.model_tester	,    '''encoder_seq_length'''	,    self.model_tester.seq_length       )
							a_		:	List[Any]									=   getattr(self.model_tester	,    '''key_length'''	,    __SCREAMING_SNAKE_CASE       )

							for model_class in self.all_model_classes:
											a_		:	str									=   self._prepare_for_class(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )
											a_		:	Optional[int]									=   model_class(__SCREAMING_SNAKE_CASE       )
											a_		:	Any									=   len(model(__SCREAMING_SNAKE_CASE       )       )

											with tempfile.TemporaryDirectory() as tmpdirname:
															model.save_pretrained(__SCREAMING_SNAKE_CASE	,    saved_model=__SCREAMING_SNAKE_CASE       )
															a_		:	Union[str, Any]									=   os.path.join(__SCREAMING_SNAKE_CASE	,    '''saved_model'''	,    '''1'''       )
															a_		:	Optional[int]									=   tf.keras.models.load_model(__SCREAMING_SNAKE_CASE       )
															a_		:	Any									=   model(__SCREAMING_SNAKE_CASE       )

															if self.is_encoder_decoder:
																			a_		:	Optional[int]									=   outputs['''encoder_hidden_states''']
																			a_		:	Dict									=   outputs['''encoder_attentions''']
															else:
																			a_		:	Union[str, Any]									=   outputs['''hidden_states''']
																			a_		:	Tuple									=   outputs['''attentions''']

															self.assertEqual(len(__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )

															a_		:	Union[str, Any]									=   getattr(
															    self.model_tester	,    '''expected_num_hidden_layers'''	,    self.model_tester.num_hidden_layers + 1       )

															self.assertEqual(len(__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
															self.assertListEqual(
															    list(output_hidden_states[0].shape[-2:]       )	,    [self.model_tester.seq_length, self.model_tester.hidden_size]	,    )

															self.assertEqual(len(__SCREAMING_SNAKE_CASE       )	,    self.model_tester.num_hidden_layers       )
															self.assertListEqual(
															    list(output_attentions[0].shape[-3:]       )	,    [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]	,    )






			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      Optional[int]:
							a_		:	Tuple									=   TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''       )
							self.assertIsNotNone(__SCREAMING_SNAKE_CASE       )







			def       SCREAMING_SNAKE_CASE   (			self   :      Any       )      ->      List[Any]:
							a_					,      a_		:	Union[str, Any]									=   self.model_tester.prepare_config_and_inputs_for_common()
							a_		:	Optional[Any]									=   True
							a_		:	Dict									=   getattr(self.model_tester	,    '''decoder_seq_length'''	,    self.model_tester.seq_length       )
							a_		:	List[str]									=   getattr(self.model_tester	,    '''encoder_seq_length'''	,    self.model_tester.seq_length       )
							a_		:	Dict									=   getattr(self.model_tester	,    '''key_length'''	,    __SCREAMING_SNAKE_CASE       )
							a_		:	Any									=   getattr(self.model_tester	,    '''key_length'''	,    __SCREAMING_SNAKE_CASE       )

							def check_decoder_attentions_output(__SCREAMING_SNAKE_CASE   :      List[str]       ):
											a_		:	int									=   len(__SCREAMING_SNAKE_CASE       )
											self.assertEqual(out_len % 2	,    0       )
											a_		:	Dict									=   outputs.decoder_attentions
											self.assertEqual(len(__SCREAMING_SNAKE_CASE       )	,    self.model_tester.num_hidden_layers       )
											self.assertListEqual(
											    list(decoder_attentions[0].shape[-3:]       )	,    [self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]	,    )

							def check_encoder_attentions_output(__SCREAMING_SNAKE_CASE   :      List[str]       ):
											a_		:	Optional[int]									=   [
											    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
											]
											self.assertEqual(len(__SCREAMING_SNAKE_CASE       )	,    self.model_tester.num_hidden_layers       )
											self.assertListEqual(
											    list(attentions[0].shape[-3:]       )	,    [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]	,    )

							for model_class in self.all_model_classes:
											a_		:	Tuple									=   True
											a_		:	Tuple									=   False
											a_		:	Optional[Any]									=   model_class(__SCREAMING_SNAKE_CASE       )
											a_		:	List[Any]									=   model(self._prepare_for_class(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )       )
											a_		:	Union[str, Any]									=   len(__SCREAMING_SNAKE_CASE       )
											self.assertEqual(config.output_hidden_states	,    __SCREAMING_SNAKE_CASE       )
											check_encoder_attentions_output(__SCREAMING_SNAKE_CASE       )

											if self.is_encoder_decoder:
															a_		:	Any									=   model_class(__SCREAMING_SNAKE_CASE       )
															a_		:	List[Any]									=   model(self._prepare_for_class(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )       )
															self.assertEqual(config.output_hidden_states	,    __SCREAMING_SNAKE_CASE       )
															check_decoder_attentions_output(__SCREAMING_SNAKE_CASE       )

											# Check that output attentions can also be changed via the config
											del inputs_dict["output_attentions"]
											a_		:	Union[str, Any]									=   True
											a_		:	Union[str, Any]									=   model_class(__SCREAMING_SNAKE_CASE       )
											a_		:	int									=   model(self._prepare_for_class(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )       )
											self.assertEqual(config.output_hidden_states	,    __SCREAMING_SNAKE_CASE       )
											check_encoder_attentions_output(__SCREAMING_SNAKE_CASE       )

											# Check attention is always last and order is fine
											a_		:	Dict									=   True
											a_		:	str									=   True
											a_		:	int									=   model_class(__SCREAMING_SNAKE_CASE       )
											a_		:	List[Any]									=   model(self._prepare_for_class(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )       )

											self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)	,    len(__SCREAMING_SNAKE_CASE       )       )
											self.assertEqual(model.config.output_hidden_states	,    __SCREAMING_SNAKE_CASE       )
											check_encoder_attentions_output(__SCREAMING_SNAKE_CASE       )





@require_tf
class 						SCREAMING_SNAKE_CASE       (       unittest.TestCase							):






			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      int       )      ->      Any:
							a_		:	int									=   TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''       )
							a_		:	List[Any]									=   tf.constant([[0, 1, 2, 3, 4, 5]]       )
							a_		:	Any									=   model(__SCREAMING_SNAKE_CASE       )[0]

							a_		:	str									=   [1, 6, 768]
							self.assertEqual(output.shape	,    __SCREAMING_SNAKE_CASE       )

							a_		:	Dict									=   tf.constant(
							    [
							        [
							            [-0.0347_5493, -0.468_6034, -0.3063_8832],
							            [0.2263_7248, -0.2698_8646, -0.742_3424],
							            [0.1032_4868, -0.4501_3508, -0.5828_0784],
							        ]
							    ]       )
							tf.debugging.assert_near(output[:, :3, :3]	,    __SCREAMING_SNAKE_CASE	,    atol=1e-4       )




