



import functools
import math

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


def       UpperCamelCase(  __UpperCamelCase							:		int					,__UpperCamelCase							:		Any					,__UpperCamelCase							:		Dict					,__UpperCamelCase							:		Optional[Any]					,__UpperCamelCase							:		int = 4096			):
		lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					Union[str, Any]     =			key.shape[-3:]
		lowerCAmelCase_      :					Any     =			value.shape[-1]
		lowerCAmelCase_      :					Any     =			min(__UpperCamelCase					,__UpperCamelCase			)
		lowerCAmelCase_      :					Dict     =			query / jnp.sqrt(__UpperCamelCase			)

		@functools.partial(jax.checkpoint					,prevent_cse=__UpperCamelCase			)
		def summarize_chunk(__UpperCamelCase							:		Optional[Any]					,__UpperCamelCase							:		int					,__UpperCamelCase							:		Dict			):
				lowerCAmelCase_      :					Optional[int]     =			jnp.einsum('''...qhd,...khd->...qhk'''					,__UpperCamelCase					,__UpperCamelCase					,precision=__UpperCamelCase			)

				lowerCAmelCase_      :					Union[str, Any]     =			jnp.max(__UpperCamelCase					,axis=-1					,keepdims=__UpperCamelCase			)
				lowerCAmelCase_      :					List[Any]     =			jax.lax.stop_gradient(__UpperCamelCase			)
				lowerCAmelCase_      :					List[str]     =			jnp.exp(attn_weights - max_score			)

				lowerCAmelCase_      :					Tuple     =			jnp.einsum('''...vhf,...qhv->...qhf'''					,__UpperCamelCase					,__UpperCamelCase					,precision=__UpperCamelCase			)
				lowerCAmelCase_      :					int     =			jnp.einsum('''...qhk->...qh'''					,__UpperCamelCase			)

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

		def chunk_scanner(__UpperCamelCase							:		Tuple			):
				# julienne key array
				lowerCAmelCase_      :					Optional[int]     =			jax.lax.dynamic_slice(
				    operand=__UpperCamelCase					,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
				lowerCAmelCase_      :					int     =			jax.lax.dynamic_slice(
				    operand=__UpperCamelCase					,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(__UpperCamelCase					,__UpperCamelCase					,__UpperCamelCase			)

		lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					Tuple     =			jax.lax.map(f=__UpperCamelCase					,xs=jnp.arange(0					,__UpperCamelCase					,__UpperCamelCase			)			)

		lowerCAmelCase_      :					List[Any]     =			jnp.max(__UpperCamelCase					,axis=0					,keepdims=__UpperCamelCase			)
		lowerCAmelCase_      :					Union[str, Any]     =			jnp.exp(chunk_max - global_max			)

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

		lowerCAmelCase_      :					str     =			chunk_values.sum(axis=0			)
		lowerCAmelCase_      :					Union[str, Any]     =			jnp.expand_dims(__UpperCamelCase					,-1			).sum(axis=0			)

		return all_values / all_weights


def       UpperCamelCase(  __UpperCamelCase							:		int					,__UpperCamelCase							:		Any					,__UpperCamelCase							:		str					,__UpperCamelCase							:		str=jax.lax.Precision.HIGHEST					,__UpperCamelCase							:		int = 1024					,__UpperCamelCase							:		int = 4096			):
		lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			query.shape[-3:]

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

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

		lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			jax.lax.scan(
		    f=__UpperCamelCase					,init=0					,xs=__UpperCamelCase					,length=math.ceil(num_q / query_chunk_size			)			)

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



class     __snake_case						(					nn.Module						):
							_a      				=			42
							_a      				=			8
							_a      				=			64
							_a      				=			0.0
							_a      				=			False
							_a      				=			jnp.floataa




							def        UpperCAmelCase__      (	self					:	Any):
									lowerCAmelCase_      :					Tuple     =			self.dim_head * self.heads
									lowerCAmelCase_      :					Any     =			self.dim_head**-0.5

									# Weights were exported with old names {to_q, to_k, to_v, to_out}
									lowerCAmelCase_      :					Any     =			nn.Dense(A_    ,    use_bias=A_    ,    dtype=self.dtype    ,    name='''to_q''')
									lowerCAmelCase_      :					int     =			nn.Dense(A_    ,    use_bias=A_    ,    dtype=self.dtype    ,    name='''to_k''')
									lowerCAmelCase_      :					Dict     =			nn.Dense(A_    ,    use_bias=A_    ,    dtype=self.dtype    ,    name='''to_v''')

									lowerCAmelCase_      :					str     =			nn.Dense(self.query_dim    ,    dtype=self.dtype    ,    name='''to_out_0''')
									lowerCAmelCase_      :					Optional[Any]     =			nn.Dropout(rate=self.dropout)




							def        UpperCAmelCase__      (	self					:	Any    ,    A_					:	List[Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					Any     =			tensor.shape
									lowerCAmelCase_      :					int     =			self.heads
									lowerCAmelCase_      :					Optional[int]     =			tensor.reshape(A_    ,    A_    ,    A_    ,    dim // head_size)
									lowerCAmelCase_      :					List[str]     =			jnp.transpose(A_    ,    (0, 2, 1, 3))
									lowerCAmelCase_      :					str     =			tensor.reshape(batch_size * head_size    ,    A_    ,    dim // head_size)
									return tensor




							def        UpperCAmelCase__      (	self					:	Optional[int]    ,    A_					:	int):
									lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[int]     =			tensor.shape
									lowerCAmelCase_      :					Union[str, Any]     =			self.heads
									lowerCAmelCase_      :					int     =			tensor.reshape(batch_size // head_size    ,    A_    ,    A_    ,    A_)
									lowerCAmelCase_      :					List[str]     =			jnp.transpose(A_    ,    (0, 2, 1, 3))
									lowerCAmelCase_      :					Any     =			tensor.reshape(batch_size // head_size    ,    A_    ,    dim * head_size)
									return tensor







							def __call__(	self					:	Dict    ,    A_					:	List[Any]    ,    A_					:	Any=None    ,    A_					:	Union[str, Any]=True):
									lowerCAmelCase_      :					Optional[int]     =			hidden_states if context is None else context

									lowerCAmelCase_      :					List[Any]     =			self.query(A_)
									lowerCAmelCase_      :					Tuple     =			self.key(A_)
									lowerCAmelCase_      :					List[Any]     =			self.value(A_)

									lowerCAmelCase_      :					str     =			self.reshape_heads_to_batch_dim(A_)
									lowerCAmelCase_      :					Any     =			self.reshape_heads_to_batch_dim(A_)
									lowerCAmelCase_      :					int     =			self.reshape_heads_to_batch_dim(A_)

									if self.use_memory_efficient_attention:
											lowerCAmelCase_      :					Dict     =			query_states.transpose(1    ,    0    ,    2)
											lowerCAmelCase_      :					List[str]     =			key_states.transpose(1    ,    0    ,    2)
											lowerCAmelCase_      :					Dict     =			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

											lowerCAmelCase_      :					Optional[Any]     =			query_states.shape[-3]
											if flatten_latent_dim % 6_4 == 0:
													lowerCAmelCase_      :					Optional[int]     =			int(flatten_latent_dim / 6_4)
											elif flatten_latent_dim % 1_6 == 0:
													lowerCAmelCase_      :					List[Any]     =			int(flatten_latent_dim / 1_6)
											elif flatten_latent_dim % 4 == 0:
													lowerCAmelCase_      :					Union[str, Any]     =			int(flatten_latent_dim / 4)
											else:
													lowerCAmelCase_      :					Any     =			int(A_)

											lowerCAmelCase_      :					str     =			jax_memory_efficient_attention(
											    A_    ,    A_    ,    A_    ,    query_chunk_size=A_    ,    key_chunk_size=4_0_9_6 * 4)

											lowerCAmelCase_      :					str     =			hidden_states.transpose(1    ,    0    ,    2)
									else:
											# compute attentions
											lowerCAmelCase_      :					Dict     =			jnp.einsum('''b i d, b j d->b i j'''    ,    A_    ,    A_)
											lowerCAmelCase_      :					Dict     =			attention_scores * self.scale
											lowerCAmelCase_      :					Optional[int]     =			nn.softmax(A_    ,    axis=2)

											# attend to values
											lowerCAmelCase_      :					Any     =			jnp.einsum('''b i j, b j d -> b i d'''    ,    A_    ,    A_)

									lowerCAmelCase_      :					Union[str, Any]     =			self.reshape_batch_dim_to_heads(A_)
									lowerCAmelCase_      :					Union[str, Any]     =			self.proj_attn(A_)
									return self.dropout_layer(A_    ,    deterministic=A_)





class     __snake_case						(					nn.Module						):
							_a      				=			42
							_a      				=			42
							_a      				=			42
							_a      				=			0.0
							_a      				=			False
							_a      				=			jnp.floataa
							_a      				=			False




							def        UpperCAmelCase__      (	self					:	Any):
									# self attention (or cross_attention if only_cross_attention is True)
									lowerCAmelCase_      :					Optional[Any]     =			FlaxAttention(
									    self.dim    ,    self.n_heads    ,    self.d_head    ,    self.dropout    ,    self.use_memory_efficient_attention    ,    dtype=self.dtype)
									# cross attention
									lowerCAmelCase_      :					Tuple     =			FlaxAttention(
									    self.dim    ,    self.n_heads    ,    self.d_head    ,    self.dropout    ,    self.use_memory_efficient_attention    ,    dtype=self.dtype)
									lowerCAmelCase_      :					Optional[Any]     =			FlaxFeedForward(dim=self.dim    ,    dropout=self.dropout    ,    dtype=self.dtype)
									lowerCAmelCase_      :					Dict     =			nn.LayerNorm(epsilon=1e-5    ,    dtype=self.dtype)
									lowerCAmelCase_      :					List[Any]     =			nn.LayerNorm(epsilon=1e-5    ,    dtype=self.dtype)
									lowerCAmelCase_      :					List[str]     =			nn.LayerNorm(epsilon=1e-5    ,    dtype=self.dtype)
									lowerCAmelCase_      :					int     =			nn.Dropout(rate=self.dropout)







							def __call__(	self					:	Optional[int]    ,    A_					:	Any    ,    A_					:	int    ,    A_					:	Any=True):
									# self attention
									lowerCAmelCase_      :					List[Any]     =			hidden_states
									if self.only_cross_attention:
											lowerCAmelCase_      :					List[str]     =			self.attna(self.norma(A_)    ,    A_    ,    deterministic=A_)
									else:
											lowerCAmelCase_      :					Optional[Any]     =			self.attna(self.norma(A_)    ,    deterministic=A_)
									lowerCAmelCase_      :					int     =			hidden_states + residual

									# cross attention
									lowerCAmelCase_      :					int     =			hidden_states
									lowerCAmelCase_      :					Optional[Any]     =			self.attna(self.norma(A_)    ,    A_    ,    deterministic=A_)
									lowerCAmelCase_      :					List[str]     =			hidden_states + residual

									# feed forward
									lowerCAmelCase_      :					Union[str, Any]     =			hidden_states
									lowerCAmelCase_      :					List[Any]     =			self.ff(self.norma(A_)    ,    deterministic=A_)
									lowerCAmelCase_      :					Optional[int]     =			hidden_states + residual

									return self.dropout_layer(A_    ,    deterministic=A_)





class     __snake_case						(					nn.Module						):
							_a      				=			42
							_a      				=			42
							_a      				=			42
							_a      				=			1
							_a      				=			0.0
							_a      				=			False
							_a      				=			False
							_a      				=			jnp.floataa
							_a      				=			False




							def        UpperCAmelCase__      (	self					:	int):
									lowerCAmelCase_      :					Any     =			nn.GroupNorm(num_groups=3_2    ,    epsilon=1e-5)

									lowerCAmelCase_      :					Dict     =			self.n_heads * self.d_head
									if self.use_linear_projection:
											lowerCAmelCase_      :					Union[str, Any]     =			nn.Dense(A_    ,    dtype=self.dtype)
									else:
											lowerCAmelCase_      :					Any     =			nn.Conv(
											    A_    ,    kernel_size=(1, 1)    ,    strides=(1, 1)    ,    padding='''VALID'''    ,    dtype=self.dtype    ,    )

									lowerCAmelCase_      :					Tuple     =			[
									    FlaxBasicTransformerBlock(
									        A_    ,    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:
											lowerCAmelCase_      :					Dict     =			nn.Dense(A_    ,    dtype=self.dtype)
									else:
											lowerCAmelCase_      :					Optional[Any]     =			nn.Conv(
											    A_    ,    kernel_size=(1, 1)    ,    strides=(1, 1)    ,    padding='''VALID'''    ,    dtype=self.dtype    ,    )

									lowerCAmelCase_      :					List[str]     =			nn.Dropout(rate=self.dropout)







							def __call__(	self					:	Tuple    ,    A_					:	List[Any]    ,    A_					:	int    ,    A_					:	Optional[Any]=True):
									lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					Any     =			hidden_states.shape
									lowerCAmelCase_      :					Dict     =			hidden_states
									lowerCAmelCase_      :					str     =			self.norm(A_)
									if self.use_linear_projection:
											lowerCAmelCase_      :					List[Any]     =			hidden_states.reshape(A_    ,    height * width    ,    A_)
											lowerCAmelCase_      :					int     =			self.proj_in(A_)
									else:
											lowerCAmelCase_      :					int     =			self.proj_in(A_)
											lowerCAmelCase_      :					List[Any]     =			hidden_states.reshape(A_    ,    height * width    ,    A_)

									for transformer_block in self.transformer_blocks:
											lowerCAmelCase_      :					Tuple     =			transformer_block(A_    ,    A_    ,    deterministic=A_)

									if self.use_linear_projection:
											lowerCAmelCase_      :					Dict     =			self.proj_out(A_)
											lowerCAmelCase_      :					Optional[int]     =			hidden_states.reshape(A_    ,    A_    ,    A_    ,    A_)
									else:
											lowerCAmelCase_      :					int     =			hidden_states.reshape(A_    ,    A_    ,    A_    ,    A_)
											lowerCAmelCase_      :					int     =			self.proj_out(A_)

									lowerCAmelCase_      :					Any     =			hidden_states + residual
									return self.dropout_layer(A_    ,    deterministic=A_)





class     __snake_case						(					nn.Module						):
							_a      				=			42
							_a      				=			0.0
							_a      				=			jnp.floataa




							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									# The second linear layer needs to be called
									# net_2 for now to match the index of the Sequential layer
									lowerCAmelCase_      :					int     =			FlaxGEGLU(self.dim    ,    self.dropout    ,    self.dtype)
									lowerCAmelCase_      :					Any     =			nn.Dense(self.dim    ,    dtype=self.dtype)







							def __call__(	self					:	Union[str, Any]    ,    A_					:	Union[str, Any]    ,    A_					:	Optional[Any]=True):
									lowerCAmelCase_      :					int     =			self.net_a(A_    ,    deterministic=A_)
									lowerCAmelCase_      :					List[str]     =			self.net_a(A_)
									return hidden_states





class     __snake_case						(					nn.Module						):
							_a      				=			42
							_a      				=			0.0
							_a      				=			jnp.floataa




							def        UpperCAmelCase__      (	self					:	int):
									lowerCAmelCase_      :					List[Any]     =			self.dim * 4
									lowerCAmelCase_      :					Any     =			nn.Dense(inner_dim * 2    ,    dtype=self.dtype)
									lowerCAmelCase_      :					Union[str, Any]     =			nn.Dropout(rate=self.dropout)







							def __call__(	self					:	Tuple    ,    A_					:	List[Any]    ,    A_					:	List[str]=True):
									lowerCAmelCase_      :					List[str]     =			self.proj(A_)
									lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			jnp.split(A_    ,    2    ,    axis=2)
									return self.dropout_layer(hidden_linear * nn.gelu(A_)    ,    deterministic=A_)










from ...configuration_utils import PretrainedConfig
from ...utils import logging


A__  :  Tuple       = logging.get_logger(__name__)

A__  :  Tuple       = {
    '''facebook/opt-125m''': '''https://huggingface.co/facebook/opt-125m/blob/main/config.json''',
    '''facebook/opt-350m''': '''https://huggingface.co/facebook/opt-350m/blob/main/config.json''',
    '''facebook/opt-1.3b''': '''https://huggingface.co/facebook/opt-1.3b/blob/main/config.json''',
    '''facebook/opt-2.7b''': '''https://huggingface.co/facebook/opt-2.7b/blob/main/config.json''',
    '''facebook/opt-6.7b''': '''https://huggingface.co/facebook/opt-6.7b/blob/main/config.json''',
    '''facebook/opt-13b''': '''https://huggingface.co/facebook/opt-13b/blob/main/config.json''',
}



class     __snake_case						(					UpperCamelCase_						):
							_a      				=			'''opt'''
							_a      				=			['''past_key_values''']




							def __init__(	self					:	Optional[int]    ,    A_					:	Any=5_0_2_7_2    ,    A_					:	Tuple=7_6_8    ,    A_					:	Optional[Any]=1_2    ,    A_					:	str=3_0_7_2    ,    A_					:	Dict=2_0_4_8    ,    A_					:	List[Any]=True    ,    A_					:	Optional[int]=False    ,    A_					:	Optional[int]=None    ,    A_					:	Dict=0.1    ,    A_					:	str=0.0    ,    A_					:	Union[str, Any]=1_2    ,    A_					:	List[str]="relu"    ,    A_					:	Dict=0.0    ,    A_					:	Union[str, Any]=0.02    ,    A_					:	Optional[Any]=True    ,    A_					:	Optional[Any]=1    ,    A_					:	List[Any]=2    ,    A_					:	Dict=2    ,    A_					:	Optional[int]=True    ,    A_					:	Optional[int]=True    ,    **A_					:	List[Any]    ,    ):
									super().__init__(
									    pad_token_id=A_    ,    bos_token_id=A_    ,    eos_token_id=A_    ,    **A_    ,    )
									lowerCAmelCase_      :					List[Any]     =			vocab_size
									lowerCAmelCase_      :					Union[str, Any]     =			max_position_embeddings
									lowerCAmelCase_      :					Optional[int]     =			num_attention_heads
									lowerCAmelCase_      :					int     =			word_embed_proj_dim if word_embed_proj_dim is not None else hidden_size
									lowerCAmelCase_      :					List[Any]     =			ffn_dim
									lowerCAmelCase_      :					Any     =			hidden_size
									lowerCAmelCase_      :					Union[str, Any]     =			num_hidden_layers
									lowerCAmelCase_      :					str     =			dropout
									lowerCAmelCase_      :					List[str]     =			attention_dropout
									lowerCAmelCase_      :					Any     =			activation_function
									lowerCAmelCase_      :					Union[str, Any]     =			init_std
									lowerCAmelCase_      :					int     =			layerdrop
									lowerCAmelCase_      :					Optional[int]     =			use_cache
									lowerCAmelCase_      :					int     =			do_layer_norm_before
									# We keep these variables at `True` for backward compatibility.
									lowerCAmelCase_      :					Optional[Any]     =			enable_bias
									lowerCAmelCase_      :					Optional[int]     =			layer_norm_elementwise_affine

									# Note that the only purpose of `_remove_final_layer_norm` is to keep backward compatibility
									# with checkpoints that have been fine-tuned before transformers v4.20.1
									# see https://github.com/facebookresearch/metaseq/pull/164
									lowerCAmelCase_      :					Optional[Any]     =			_remove_final_layer_norm





