



"""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     :


					'''simple docstring'''





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



					def 		snake_case	(	self							:  int					)-> Tuple:
								lowerCamelCase__		:		Optional[Any]		     =ids_tensor([self.batch_size, self.seq_length],  self.vocab_size					)

								lowerCamelCase__		:		Optional[int]		     =None
								if self.use_input_mask:
											lowerCamelCase__		:		str		     =random_attention_mask([self.batch_size, self.seq_length]					)

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

								lowerCamelCase__		:		Tuple		     =None
								lowerCamelCase__		:		List[Any]		     =None
								lowerCamelCase__		:		Tuple		     =None
								if self.use_labels:
											lowerCamelCase__		:		int		     =ids_tensor([self.batch_size],  self.type_sequence_label_size					)
											lowerCamelCase__		:		Dict		     =ids_tensor([self.batch_size, self.seq_length],  self.num_labels					)
											lowerCamelCase__		:		Tuple		     =ids_tensor([self.batch_size],  self.num_choices					)

								lowerCamelCase__		:		List[str]		     =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=lowerCamelCase,  )

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



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

								lowerCamelCase__		:		int		     =[input_ids, input_mask]
								lowerCamelCase__		:		Any		     =model(lowerCamelCase					)

								lowerCamelCase__		:		Optional[Any]		     =model(lowerCamelCase					)

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



					def 		snake_case	(	self							:  Tuple,  lowerCamelCase							:  str,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Any,  lowerCamelCase							:  Any,  lowerCamelCase							:  List[Any],  lowerCamelCase							:  Any,  lowerCamelCase							:  Optional[Any]					)-> Union[str, Any]:
								lowerCamelCase__		:		Union[str, Any]		     =TFConvBertForMaskedLM(config=lowerCamelCase					)
								lowerCamelCase__		:		Union[str, Any]		     ={
								    'input_ids': input_ids,
								    'attention_mask': input_mask,
								    'token_type_ids': token_type_ids,
								}
								lowerCamelCase__		:		int		     =model(lowerCamelCase					)
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.seq_length, self.vocab_size)					)



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

								lowerCamelCase__		:		Dict		     =model(lowerCamelCase					)
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels)					)



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



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



					def 		snake_case	(	self							:  int,  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Dict,  lowerCamelCase							:  List[Any],  lowerCamelCase							:  List[Any],  lowerCamelCase							:  Tuple					)-> Optional[int]:
								lowerCamelCase__		:		str		     =TFConvBertForQuestionAnswering(config=lowerCamelCase					)
								lowerCamelCase__		:		Dict		     ={
								    'input_ids': input_ids,
								    'attention_mask': input_mask,
								    'token_type_ids': token_type_ids,
								}

								lowerCamelCase__		:		str		     =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							:  Tuple					)-> Any:
								lowerCamelCase__		:		Union[str, Any]		     =self.prepare_config_and_inputs()
								(
								    lowerCamelCase__    
								)		:		Union[str, Any]		     =config_and_inputs
								lowerCamelCase__		:		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     (				_a		,       _a		,       unittest.TestCase       ):


					'''simple docstring'''





					_a             =							(
					    (
					        TFConvBertModel,
					        TFConvBertForMaskedLM,
					        TFConvBertForQuestionAnswering,
					        TFConvBertForSequenceClassification,
					        TFConvBertForTokenClassification,
					        TFConvBertForMultipleChoice,
					    )
					    if is_tf_available()
					    else ()
					)
					_a             =							(
					    {
					        """feature-extraction""": TFConvBertModel,
					        """fill-mask""": TFConvBertForMaskedLM,
					        """question-answering""": TFConvBertForQuestionAnswering,
					        """text-classification""": TFConvBertForSequenceClassification,
					        """token-classification""": TFConvBertForTokenClassification,
					        """zero-shot""": TFConvBertForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					_a             =							False
					_a             =							False
					_a             =							False



					def 		snake_case	(	self							:  Optional[int]					)-> Optional[int]:
								lowerCamelCase__		:		List[str]		     =TFConvBertModelTester(self					)
								lowerCamelCase__		:		List[Any]		     =ConfigTester(self,  config_class=lowerCamelCase,  hidden_size=37					)



					def 		snake_case	(	self							:  Dict					)-> List[str]:
								self.config_tester.run_common_tests()



					def 		snake_case	(	self							:  Dict					)-> Union[str, Any]:
								lowerCamelCase__		:		Union[str, Any]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  Optional[int]					)-> List[str]:
								lowerCamelCase__		:		int		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase					)



					def 		snake_case	(	self							:  Optional[int]					)-> Optional[Any]:
								lowerCamelCase__		:		Tuple		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase					)



					def 		snake_case	(	self							:  int					)-> Optional[Any]:
								lowerCamelCase__		:		str		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*lowerCamelCase					)



					def 		snake_case	(	self							:  List[str]					)-> Dict:
								lowerCamelCase__		:		List[str]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase					)



					def 		snake_case	(	self							:  str					)-> int:
								lowerCamelCase__		:		Optional[Any]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*lowerCamelCase					)



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

								if hasattr(lowerCamelCase,  '''use_cache'''					):
											lowerCamelCase__		:		str		     =True

								lowerCamelCase__		:		Optional[Any]		     =getattr(self.model_tester,  '''encoder_seq_length''',  self.model_tester.seq_length					)
								lowerCamelCase__		:		str		     =getattr(self.model_tester,  '''key_length''',  lowerCamelCase					)

								for model_class in self.all_model_classes:
											lowerCamelCase__		:		Union[str, Any]		     =self._prepare_for_class(lowerCamelCase,  lowerCamelCase					)
											lowerCamelCase__		:		Tuple		     =model_class(lowerCamelCase					)
											lowerCamelCase__		:		Optional[Any]		     =len(model(lowerCamelCase					)					)

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

														if self.is_encoder_decoder:
																	lowerCamelCase__		:		Tuple		     =outputs['encoder_hidden_states']
																	lowerCamelCase__		:		Tuple		     =outputs['encoder_attentions']
														else:
																	lowerCamelCase__		:		List[str]		     =outputs['hidden_states']
																	lowerCamelCase__		:		Dict		     =outputs['attentions']

														self.assertEqual(len(lowerCamelCase					),  lowerCamelCase					)

														lowerCamelCase__		:		Tuple		     =getattr(
														    self.model_tester,  '''expected_num_hidden_layers''',  self.model_tester.num_hidden_layers + 1					)

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

														self.assertEqual(len(lowerCamelCase					),  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 		snake_case	(	self							:  List[str]					)-> Optional[int]:
								lowerCamelCase__		:		Optional[int]		     =TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''					)
								self.assertIsNotNone(lowerCamelCase					)




					def 		snake_case	(	self							:  int					)-> Dict:
								lowerCamelCase__		:		List[str]		     =self.model_tester.prepare_config_and_inputs_for_common()
								lowerCamelCase__		:		int		     =True
								lowerCamelCase__		:		Optional[int]		     =getattr(self.model_tester,  '''decoder_seq_length''',  self.model_tester.seq_length					)
								lowerCamelCase__		:		Union[str, Any]		     =getattr(self.model_tester,  '''encoder_seq_length''',  self.model_tester.seq_length					)
								lowerCamelCase__		:		Union[str, Any]		     =getattr(self.model_tester,  '''key_length''',  lowerCamelCase					)
								lowerCamelCase__		:		Union[str, Any]		     =getattr(self.model_tester,  '''key_length''',  lowerCamelCase					)

								def check_decoder_attentions_output(lowerCamelCase							:  int					):
											lowerCamelCase__		:		List[str]		     =len(lowerCamelCase					)
											self.assertEqual(out_len % 2,  0					)
											lowerCamelCase__		:		Union[str, Any]		     =outputs.decoder_attentions
											self.assertEqual(len(lowerCamelCase					),  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(lowerCamelCase							:  int					):
											lowerCamelCase__		:		Dict		     =[
											    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
											]
											self.assertEqual(len(lowerCamelCase					),  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:
											lowerCamelCase__		:		int		     =True
											lowerCamelCase__		:		Tuple		     =False
											lowerCamelCase__		:		str		     =model_class(lowerCamelCase					)
											lowerCamelCase__		:		str		     =model(self._prepare_for_class(lowerCamelCase,  lowerCamelCase					)					)
											lowerCamelCase__		:		int		     =len(lowerCamelCase					)
											self.assertEqual(config.output_hidden_states,  lowerCamelCase					)
											check_encoder_attentions_output(lowerCamelCase					)

											if self.is_encoder_decoder:
														lowerCamelCase__		:		str		     =model_class(lowerCamelCase					)
														lowerCamelCase__		:		str		     =model(self._prepare_for_class(lowerCamelCase,  lowerCamelCase					)					)
														self.assertEqual(config.output_hidden_states,  lowerCamelCase					)
														check_decoder_attentions_output(lowerCamelCase					)

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

											# Check attention is always last and order is fine
											lowerCamelCase__		:		Tuple		     =True
											lowerCamelCase__		:		List[Any]		     =True
											lowerCamelCase__		:		Optional[int]		     =model_class(lowerCamelCase					)
											lowerCamelCase__		:		List[str]		     =model(self._prepare_for_class(lowerCamelCase,  lowerCamelCase					)					)

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





@require_tf
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					@slow
					def 		snake_case	(	self							:  Tuple					)-> List[Any]:
								lowerCamelCase__		:		Union[str, Any]		     =TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''					)
								lowerCamelCase__		:		List[Any]		     =tf.constant([[0, 1, 2, 3, 4, 5]]					)
								lowerCamelCase__		:		Tuple		     =model(lowerCamelCase					)[0]

								lowerCamelCase__		:		List[str]		     =[1, 6, 768]
								self.assertEqual(output.shape,  lowerCamelCase					)

								lowerCamelCase__		:		List[str]		     =tf.constant(
								    [
								        [
								            [-0.03_475_493, -0.4_686_034, -0.30_638_832],
								            [0.22_637_248, -0.26_988_646, -0.7_423_424],
								            [0.10_324_868, -0.45_013_508, -0.58_280_784],
								        ]
								    ]					)
								tf.debugging.assert_near(output[:, :3, :3],  lowerCamelCase,  atol=1E-4					)









"""simple docstring"""






import math

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






def 					snake_case__       (		__lowerCamelCase    :   jnp.ndarray							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1.0e4							,		__lowerCamelCase    :   bool = False							,		__lowerCamelCase    :   float = 1.0							,		):
			"""simple docstring"""
			assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
			assert embedding_dim % 2 == 0, f'''Embedding dimension {embedding_dim} should be even'''
			lowerCamelCase__		:		Any		     =float(embedding_dim // 2     )
			lowerCamelCase__		:		List[str]		     =math.log(max_timescale / min_timescale     ) / (num_timescales - freq_shift)
			lowerCamelCase__		:		int		     =min_timescale * jnp.exp(jnp.arange(__lowerCamelCase							,		dtype=jnp.floataa     ) * -log_timescale_increment     )
			lowerCamelCase__		:		Tuple		     =jnp.expand_dims(__lowerCamelCase							,		1     ) * jnp.expand_dims(__lowerCamelCase							,		0     )

			# scale embeddings
			lowerCamelCase__		:		List[str]		     =scale * emb

			if flip_sin_to_cos:
						lowerCamelCase__		:		int		     =jnp.concatenate([jnp.cos(__lowerCamelCase     ), jnp.sin(__lowerCamelCase     )]							,		axis=1     )
			else:
						lowerCamelCase__		:		List[str]		     =jnp.concatenate([jnp.sin(__lowerCamelCase     ), jnp.cos(__lowerCamelCase     )]							,		axis=1     )
			lowerCamelCase__		:		str		     =jnp.reshape(__lowerCamelCase							,		[jnp.shape(__lowerCamelCase     )[0], embedding_dim]     )
			return signal

class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							jnp.floataa



					@nn.compact
					def __call__(	self							:  Optional[Any],  lowerCamelCase							:  int					)-> Any:
								lowerCamelCase__		:		Optional[Any]		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_1'''					)(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =nn.silu(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_2'''					)(lowerCamelCase					)
								return temb





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							False
					_a             =							1



					@nn.compact
					def __call__(	self							:  Any,  lowerCamelCase							:  int					)-> int:
								return get_sinusoidal_embeddings(
								    lowerCamelCase,  embedding_dim=self.dim,  flip_sin_to_cos=self.flip_sin_to_cos,  freq_shift=self.freq_shift					)




