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






		"""simple docstring"""







		def __init__(			self    :				Dict  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				Tuple=13  ,					__SCREAMING_SNAKE_CASE    :				List[str]=7  ,					__SCREAMING_SNAKE_CASE    :				Dict=True  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]=True  ,					__SCREAMING_SNAKE_CASE    :				int=True  ,					__SCREAMING_SNAKE_CASE    :				Dict=True  ,					__SCREAMING_SNAKE_CASE    :				int=99  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]=16  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]=32  ,					__SCREAMING_SNAKE_CASE    :				Dict=2  ,					__SCREAMING_SNAKE_CASE    :				Tuple=4  ,					__SCREAMING_SNAKE_CASE    :				Any=37  ,					__SCREAMING_SNAKE_CASE    :				Dict="gelu"  ,					__SCREAMING_SNAKE_CASE    :				int=0.1  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=0.1  ,					__SCREAMING_SNAKE_CASE    :				Tuple=512  ,					__SCREAMING_SNAKE_CASE    :				int=16  ,					__SCREAMING_SNAKE_CASE    :				str=2  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]=0.0_2  ,					__SCREAMING_SNAKE_CASE    :				List[str]=3  ,					__SCREAMING_SNAKE_CASE    :				Any=4  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=None  ,					):






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



		def 				_a				(			self    :				Any   ):






					"""simple docstring"""
					_snake_case              =      ids_tensor([self.batch_size, self.seq_length]  ,					self.vocab_size   )

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

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

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

					_snake_case              =      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 				_a				(			self    :				Any  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				List[Any]   ):






					"""simple docstring"""
					_snake_case              =      TFAlbertModel(config=__SCREAMING_SNAKE_CASE   )
					# inputs = {'input_ids': input_ids,
					#           'attention_mask': input_mask,
					#           'token_type_ids': token_type_ids}
					# sequence_output, pooled_output = model(**inputs)
					_snake_case              =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					_snake_case              =      model(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      [input_ids, input_mask]
					_snake_case              =      model(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      model(__SCREAMING_SNAKE_CASE   )

					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 				_a				(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				Dict  ,					__SCREAMING_SNAKE_CASE    :				Dict   ):






					"""simple docstring"""
					_snake_case              =      self.num_labels
					_snake_case              =      TFAlbertForPreTraining(config=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					_snake_case              =      model(__SCREAMING_SNAKE_CASE   )
					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 				_a				(			self    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]   ):






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



		def 				_a				(			self    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				Optional[Any]  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]   ):






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



		def 				_a				(			self    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				str   ):






					"""simple docstring"""
					_snake_case              =      TFAlbertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					_snake_case              =      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 				_a				(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				str  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Dict  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]   ):






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



		def 				_a				(			self    :				str  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				Dict  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				List[Any]   ):






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







		def 				_a				(			self    :				Any   ):






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

@require_tf
class 	_SCREAMING_SNAKE_CASE							(     _lowercase      ,     _lowercase      ,     unittest.TestCase							):






		"""simple docstring"""







		A_						 =      (
		    (
		        TFAlbertModel,
		        TFAlbertForPreTraining,
		        TFAlbertForMaskedLM,
		        TFAlbertForSequenceClassification,
		        TFAlbertForQuestionAnswering,
		        TFAlbertForTokenClassification,
		        TFAlbertForMultipleChoice,
		    )
		    if is_tf_available()
		    else ()
		)
		A_						 =      (
		    {
		        '''feature-extraction''': TFAlbertModel,
		        '''fill-mask''': TFAlbertForMaskedLM,
		        '''question-answering''': TFAlbertForQuestionAnswering,
		        '''text-classification''': TFAlbertForSequenceClassification,
		        '''token-classification''': TFAlbertForTokenClassification,
		        '''zero-shot''': TFAlbertForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		A_						 =      False
		A_						 =      False



		def 				_a				(			self    :				int  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				List[str]=False   ):






					"""simple docstring"""
					_snake_case              =      super()._prepare_for_class(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					return_labels=__SCREAMING_SNAKE_CASE   )

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

					return inputs_dict



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      TFAlbertModelTester(self   )
					_snake_case              =      ConfigTester(self  ,					config_class=__SCREAMING_SNAKE_CASE  ,					hidden_size=37   )



		def 				_a				(			self    :				str   ):






					"""simple docstring"""
					self.config_tester.run_common_tests()



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_model(*__SCREAMING_SNAKE_CASE   )



		def 				_a				(			self    :				Optional[Any]   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_pretraining(*__SCREAMING_SNAKE_CASE   )



		def 				_a				(			self    :				List[str]   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_masked_lm(*__SCREAMING_SNAKE_CASE   )



		def 				_a				(			self    :				str   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_multiple_choice(*__SCREAMING_SNAKE_CASE   )



		def 				_a				(			self    :				List[Any]   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_sequence_classification(*__SCREAMING_SNAKE_CASE   )



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_question_answering(*__SCREAMING_SNAKE_CASE   )







		@slow
		def 				_a				(			self    :				Any   ):






					"""simple docstring"""
					for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								_snake_case              =      TFAlbertModel.from_pretrained(__SCREAMING_SNAKE_CASE   )
								self.assertIsNotNone(__SCREAMING_SNAKE_CASE   )

@require_tf
class 	_SCREAMING_SNAKE_CASE							(     unittest.TestCase							):






		"""simple docstring"""







		@slow
		def 				_a				(			self    :				List[Any]   ):






					"""simple docstring"""
					_snake_case              =      TFAlbertForPreTraining.from_pretrained("albert-base-v2"   )
					_snake_case              =      tf.constant([[0, 1, 2, 3, 4, 5]]   )
					_snake_case              =      model(__SCREAMING_SNAKE_CASE   )[0]

					_snake_case              =      [1, 6, 30_000]
					self.assertEqual(output.shape  ,					__SCREAMING_SNAKE_CASE   )

					_snake_case              =      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]  ,					__SCREAMING_SNAKE_CASE  ,					atol=1e-4   )

import functools
import math

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






def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_ = 4096  ):



			"""simple docstring"""





			_snake_case       , _snake_case       , _snake_case              =      key.shape[-3:]
			_snake_case              =      value.shape[-1]
			_snake_case              =      min(lowerCamelCase_							,	lowerCamelCase_  )
			_snake_case              =      query / jnp.sqrt(lowerCamelCase_  )

			@functools.partial(jax.checkpoint							,	prevent_cse=lowerCamelCase_  )
			def summarize_chunk(lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  ):
						_snake_case              =      jnp.einsum("...qhd,...khd->...qhk"							,	lowerCamelCase_							,	lowerCamelCase_							,	precision=lowerCamelCase_  )

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

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

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

			def chunk_scanner(lowerCamelCase_  ):
						# julienne key array
						_snake_case              =      jax.lax.dynamic_slice(
						    operand=lowerCamelCase_							,	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
						_snake_case              =      jax.lax.dynamic_slice(
						    operand=lowerCamelCase_							,	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(lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_  )

			_snake_case       , _snake_case       , _snake_case              =      jax.lax.map(f=lowerCamelCase_							,	xs=jnp.arange(0							,	lowerCamelCase_							,	lowerCamelCase_  )  )

			_snake_case              =      jnp.max(lowerCamelCase_							,	axis=0							,	keepdims=lowerCamelCase_  )
			_snake_case              =      jnp.exp(chunk_max - global_max  )

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

			_snake_case              =      chunk_values.sum(axis=0  )
			_snake_case              =      jnp.expand_dims(lowerCamelCase_							,	-1  ).sum(axis=0  )

			return all_values / all_weights






def    UpperCAmelCase__(						lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_							,	lowerCamelCase_=jax.lax.Precision.HIGHEST							,	lowerCamelCase_ = 1024							,	lowerCamelCase_ = 4096  ):



			"""simple docstring"""





			_snake_case       , _snake_case       , _snake_case              =      query.shape[-3:]

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

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

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

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

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      8
		A_						 =      64
		A_						 =      0.0
		A_						 =      False
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      self.dim_head * self.heads
					_snake_case              =      self.dim_head**-0.5

					# Weights were exported with old names {to_q, to_k, to_v, to_out}
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_q"   )
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_k"   )
					_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					use_bias=__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype  ,					name="to_v"   )

					_snake_case              =      nn.Dense(self.query_dim  ,					dtype=self.dtype  ,					name="to_out_0"   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )



		def 				_a				(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Any   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case              =      tensor.shape
					_snake_case              =      self.heads
					_snake_case              =      tensor.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					dim // head_size   )
					_snake_case              =      jnp.transpose(__SCREAMING_SNAKE_CASE  ,					(0, 2, 1, 3)   )
					_snake_case              =      tensor.reshape(batch_size * head_size  ,					__SCREAMING_SNAKE_CASE  ,					dim // head_size   )
					return tensor



		def 				_a				(			self    :				str  ,					__SCREAMING_SNAKE_CASE    :				Tuple   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case              =      tensor.shape
					_snake_case              =      self.heads
					_snake_case              =      tensor.reshape(batch_size // head_size  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
					_snake_case              =      jnp.transpose(__SCREAMING_SNAKE_CASE  ,					(0, 2, 1, 3)   )
					_snake_case              =      tensor.reshape(batch_size // head_size  ,					__SCREAMING_SNAKE_CASE  ,					dim * head_size   )
					return tensor







		def __call__(			self    :				int  ,					__SCREAMING_SNAKE_CASE    :				int  ,					__SCREAMING_SNAKE_CASE    :				Tuple=None  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=True   ):






					"""simple docstring"""
					_snake_case              =      hidden_states if context is None else context

					_snake_case              =      self.query(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.key(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.value(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.reshape_heads_to_batch_dim(__SCREAMING_SNAKE_CASE   )

					if self.use_memory_efficient_attention:
								_snake_case              =      query_states.transpose(1  ,					0  ,					2   )
								_snake_case              =      key_states.transpose(1  ,					0  ,					2   )
								_snake_case              =      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

								_snake_case              =      query_states.shape[-3]
								if flatten_latent_dim % 64 == 0:
											_snake_case              =      int(flatten_latent_dim / 64   )
								elif flatten_latent_dim % 16 == 0:
											_snake_case              =      int(flatten_latent_dim / 16   )
								elif flatten_latent_dim % 4 == 0:
											_snake_case              =      int(flatten_latent_dim / 4   )
								else:
											_snake_case              =      int(__SCREAMING_SNAKE_CASE   )

								_snake_case              =      jax_memory_efficient_attention(
								    __SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					query_chunk_size=__SCREAMING_SNAKE_CASE  ,					key_chunk_size=4_096 * 4   )

								_snake_case              =      hidden_states.transpose(1  ,					0  ,					2   )
					else:
								# compute attentions
								_snake_case              =      jnp.einsum("b i d, b j d->b i j"  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      attention_scores * self.scale
								_snake_case              =      nn.softmax(__SCREAMING_SNAKE_CASE  ,					axis=2   )

								# attend to values
								_snake_case              =      jnp.einsum("b i j, b j d -> b i d"  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )

					_snake_case              =      self.reshape_batch_dim_to_heads(__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.proj_attn(__SCREAMING_SNAKE_CASE   )
					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      42
		A_						 =      42
		A_						 =      0.0
		A_						 =      False
		A_						 =      jnp.floataa
		A_						 =      False



		def 				_a				(			self    :				Any   ):






					"""simple docstring"""
					# self attention (or cross_attention if only_cross_attention is True)
					_snake_case              =      FlaxAttention(
					    self.dim  ,					self.n_heads  ,					self.d_head  ,					self.dropout  ,					self.use_memory_efficient_attention  ,					dtype=self.dtype   )
					# cross attention
					_snake_case              =      FlaxAttention(
					    self.dim  ,					self.n_heads  ,					self.d_head  ,					self.dropout  ,					self.use_memory_efficient_attention  ,					dtype=self.dtype   )
					_snake_case              =      FlaxFeedForward(dim=self.dim  ,					dropout=self.dropout  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.LayerNorm(epsilon=1e-5  ,					dtype=self.dtype   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				List[Any]  ,					__SCREAMING_SNAKE_CASE    :				Any  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				Any=True   ):






					"""simple docstring"""
					# self attention
					_snake_case              =      hidden_states
					if self.only_cross_attention:
								_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					# cross attention
					_snake_case              =      hidden_states
					_snake_case              =      self.attna(self.norma(__SCREAMING_SNAKE_CASE   )  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					# feed forward
					_snake_case              =      hidden_states
					_snake_case              =      self.ff(self.norma(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      hidden_states + residual

					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      42
		A_						 =      42
		A_						 =      1
		A_						 =      0.0
		A_						 =      False
		A_						 =      False
		A_						 =      jnp.floataa
		A_						 =      False



		def 				_a				(			self    :				Union[str, Any]   ):






					"""simple docstring"""
					_snake_case              =      nn.GroupNorm(num_groups=32  ,					epsilon=1e-5   )

					_snake_case              =      self.n_heads * self.d_head
					if self.use_linear_projection:
								_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype   )
					else:
								_snake_case              =      nn.Conv(
								    __SCREAMING_SNAKE_CASE  ,					kernel_size=(1, 1)  ,					strides=(1, 1)  ,					padding="VALID"  ,					dtype=self.dtype  ,					)

					_snake_case              =      [
					    FlaxBasicTransformerBlock(
					        __SCREAMING_SNAKE_CASE  ,					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:
								_snake_case              =      nn.Dense(__SCREAMING_SNAKE_CASE  ,					dtype=self.dtype   )
					else:
								_snake_case              =      nn.Conv(
								    __SCREAMING_SNAKE_CASE  ,					kernel_size=(1, 1)  ,					strides=(1, 1)  ,					padding="VALID"  ,					dtype=self.dtype  ,					)

					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]=True   ):






					"""simple docstring"""
					_snake_case       , _snake_case       , _snake_case       , _snake_case              =      hidden_states.shape
					_snake_case              =      hidden_states
					_snake_case              =      self.norm(__SCREAMING_SNAKE_CASE   )
					if self.use_linear_projection:
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					height * width  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      self.proj_in(__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      self.proj_in(__SCREAMING_SNAKE_CASE   )
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					height * width  ,					__SCREAMING_SNAKE_CASE   )

					for transformer_block in self.transformer_blocks:
								_snake_case              =      transformer_block(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

					if self.use_linear_projection:
								_snake_case              =      self.proj_out(__SCREAMING_SNAKE_CASE   )
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
					else:
								_snake_case              =      hidden_states.reshape(__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE  ,					__SCREAMING_SNAKE_CASE   )
								_snake_case              =      self.proj_out(__SCREAMING_SNAKE_CASE   )

					_snake_case              =      hidden_states + residual
					return self.dropout_layer(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      0.0
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Optional[int]   ):






					"""simple docstring"""
					# The second linear layer needs to be called
					# net_2 for now to match the index of the Sequential layer
					_snake_case              =      FlaxGEGLU(self.dim  ,					self.dropout  ,					self.dtype   )
					_snake_case              =      nn.Dense(self.dim  ,					dtype=self.dtype   )







		def __call__(			self    :				Tuple  ,					__SCREAMING_SNAKE_CASE    :				List[str]  ,					__SCREAMING_SNAKE_CASE    :				List[Any]=True   ):






					"""simple docstring"""
					_snake_case              =      self.net_a(__SCREAMING_SNAKE_CASE  ,					deterministic=__SCREAMING_SNAKE_CASE   )
					_snake_case              =      self.net_a(__SCREAMING_SNAKE_CASE   )
					return hidden_states

class 	_SCREAMING_SNAKE_CASE							(     nn.Module							):






		"""simple docstring"""







		A_						 =      42
		A_						 =      0.0
		A_						 =      jnp.floataa



		def 				_a				(			self    :				Optional[Any]   ):






					"""simple docstring"""
					_snake_case              =      self.dim * 4
					_snake_case              =      nn.Dense(inner_dim * 2  ,					dtype=self.dtype   )
					_snake_case              =      nn.Dropout(rate=self.dropout   )







		def __call__(			self    :				Optional[int]  ,					__SCREAMING_SNAKE_CASE    :				Union[str, Any]  ,					__SCREAMING_SNAKE_CASE    :				List[str]=True   ):






					"""simple docstring"""
					_snake_case              =      self.proj(__SCREAMING_SNAKE_CASE   )
					_snake_case       , _snake_case              =      jnp.split(__SCREAMING_SNAKE_CASE  ,					2  ,					axis=2   )
					return self.dropout_layer(hidden_linear * nn.gelu(__SCREAMING_SNAKE_CASE   )  ,					deterministic=__SCREAMING_SNAKE_CASE   )

