


'''simple docstring'''


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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD





class        lowercase__    (		nn.Module   ):





							'''simple docstring'''


							A_				: int
							A_				: int
							A_				: float       =    0.0
							A_				: int       =    1
							A_				: int       =    1
							A_				: bool       =    True
							A_				: bool       =    False
							A_				: bool       =    False
							A_				: bool       =    False
							A_				: jnp.dtype       =    jnp.floataa






							def UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			Tuple							= []
											_SCREAMING_SNAKE_CASE					:			str							= []

											for i in range(self.num_layers    ):
															_SCREAMING_SNAKE_CASE					:			List[Any]							= self.in_channels if i == 0 else self.out_channels

															_SCREAMING_SNAKE_CASE					:			Optional[int]							= FlaxResnetBlockaD(
															    in_channels=__snake_case       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
															resnets.append(__snake_case    )

															_SCREAMING_SNAKE_CASE					:			int							= FlaxTransformeraDModel(
															    in_channels=self.out_channels       ,   n_heads=self.num_attention_heads       ,   d_head=self.out_channels // self.num_attention_heads       ,   depth=1       ,   use_linear_projection=self.use_linear_projection       ,   only_cross_attention=self.only_cross_attention       ,   use_memory_efficient_attention=self.use_memory_efficient_attention       ,   dtype=self.dtype       ,   )
															attentions.append(__snake_case    )

											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= resnets
											_SCREAMING_SNAKE_CASE					:			Dict							= attentions

											if self.add_downsample:
															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= FlaxDownsampleaD(self.out_channels       ,   dtype=self.dtype    )






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case=True    ):
											_SCREAMING_SNAKE_CASE					:			List[Any]							= ()

											for resnet, attn in zip(self.resnets       ,   self.attentions    ):
															_SCREAMING_SNAKE_CASE					:			Tuple							= resnet(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )
															_SCREAMING_SNAKE_CASE					:			Tuple							= attn(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )
															output_states += (hidden_states,)

											if self.add_downsample:
															_SCREAMING_SNAKE_CASE					:			int							= self.downsamplers_a(__snake_case    )
															output_states += (hidden_states,)

											return hidden_states, output_states





class        lowercase__    (		nn.Module   ):





							'''simple docstring'''


							A_				: int
							A_				: int
							A_				: float       =    0.0
							A_				: int       =    1
							A_				: bool       =    True
							A_				: jnp.dtype       =    jnp.floataa






							def UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			Dict							= []

											for i in range(self.num_layers    ):
															_SCREAMING_SNAKE_CASE					:			Tuple							= self.in_channels if i == 0 else self.out_channels

															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= FlaxResnetBlockaD(
															    in_channels=__snake_case       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
															resnets.append(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= resnets

											if self.add_downsample:
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= FlaxDownsampleaD(self.out_channels       ,   dtype=self.dtype    )






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case=True    ):
											_SCREAMING_SNAKE_CASE					:			Tuple							= ()

											for resnet in self.resnets:
															_SCREAMING_SNAKE_CASE					:			List[str]							= resnet(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )
															output_states += (hidden_states,)

											if self.add_downsample:
															_SCREAMING_SNAKE_CASE					:			str							= self.downsamplers_a(__snake_case    )
															output_states += (hidden_states,)

											return hidden_states, output_states





class        lowercase__    (		nn.Module   ):





							'''simple docstring'''


							A_				: int
							A_				: int
							A_				: int
							A_				: float       =    0.0
							A_				: int       =    1
							A_				: int       =    1
							A_				: bool       =    True
							A_				: bool       =    False
							A_				: bool       =    False
							A_				: bool       =    False
							A_				: jnp.dtype       =    jnp.floataa






							def UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			Dict							= []
											_SCREAMING_SNAKE_CASE					:			int							= []

											for i in range(self.num_layers    ):
															_SCREAMING_SNAKE_CASE					:			Optional[int]							= self.in_channels if (i == self.num_layers - 1) else self.out_channels
															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= self.prev_output_channel if i == 0 else self.out_channels

															_SCREAMING_SNAKE_CASE					:			int							= FlaxResnetBlockaD(
															    in_channels=resnet_in_channels + res_skip_channels       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
															resnets.append(__snake_case    )

															_SCREAMING_SNAKE_CASE					:			Any							= FlaxTransformeraDModel(
															    in_channels=self.out_channels       ,   n_heads=self.num_attention_heads       ,   d_head=self.out_channels // self.num_attention_heads       ,   depth=1       ,   use_linear_projection=self.use_linear_projection       ,   only_cross_attention=self.only_cross_attention       ,   use_memory_efficient_attention=self.use_memory_efficient_attention       ,   dtype=self.dtype       ,   )
															attentions.append(__snake_case    )

											_SCREAMING_SNAKE_CASE					:			List[Any]							= resnets
											_SCREAMING_SNAKE_CASE					:			List[Any]							= attentions

											if self.add_upsample:
															_SCREAMING_SNAKE_CASE					:			str							= FlaxUpsampleaD(self.out_channels       ,   dtype=self.dtype    )






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case=True    ):
											for resnet, attn in zip(self.resnets       ,   self.attentions    ):
															# pop res hidden states
															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= res_hidden_states_tuple[-1]
															_SCREAMING_SNAKE_CASE					:			Tuple							= res_hidden_states_tuple[:-1]
															_SCREAMING_SNAKE_CASE					:			List[str]							= jnp.concatenate((hidden_states, res_hidden_states)       ,   axis=-1    )

															_SCREAMING_SNAKE_CASE					:			Tuple							= resnet(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )
															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= attn(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )

											if self.add_upsample:
															_SCREAMING_SNAKE_CASE					:			List[Any]							= self.upsamplers_a(__snake_case    )

											return hidden_states





class        lowercase__    (		nn.Module   ):





							'''simple docstring'''


							A_				: int
							A_				: int
							A_				: int
							A_				: float       =    0.0
							A_				: int       =    1
							A_				: bool       =    True
							A_				: jnp.dtype       =    jnp.floataa






							def UpperCAmelCase_      (       self    ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= []

											for i in range(self.num_layers    ):
															_SCREAMING_SNAKE_CASE					:			Any							= self.in_channels if (i == self.num_layers - 1) else self.out_channels
															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= self.prev_output_channel if i == 0 else self.out_channels

															_SCREAMING_SNAKE_CASE					:			Dict							= FlaxResnetBlockaD(
															    in_channels=resnet_in_channels + res_skip_channels       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
															resnets.append(__snake_case    )

											_SCREAMING_SNAKE_CASE					:			List[str]							= resnets

											if self.add_upsample:
															_SCREAMING_SNAKE_CASE					:			List[str]							= FlaxUpsampleaD(self.out_channels       ,   dtype=self.dtype    )






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case=True    ):
											for resnet in self.resnets:
															# pop res hidden states
															_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= res_hidden_states_tuple[-1]
															_SCREAMING_SNAKE_CASE					:			Dict							= res_hidden_states_tuple[:-1]
															_SCREAMING_SNAKE_CASE					:			List[Any]							= jnp.concatenate((hidden_states, res_hidden_states)       ,   axis=-1    )

															_SCREAMING_SNAKE_CASE					:			Optional[Any]							= resnet(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )

											if self.add_upsample:
															_SCREAMING_SNAKE_CASE					:			Tuple							= self.upsamplers_a(__snake_case    )

											return hidden_states





class        lowercase__    (		nn.Module   ):





							'''simple docstring'''


							A_				: int
							A_				: float       =    0.0
							A_				: int       =    1
							A_				: int       =    1
							A_				: bool       =    False
							A_				: bool       =    False
							A_				: jnp.dtype       =    jnp.floataa






							def UpperCAmelCase_      (       self    ):
											# there is always at least one resnet
											_SCREAMING_SNAKE_CASE					:			Dict							= [
											    FlaxResnetBlockaD(
											        in_channels=self.in_channels       ,   out_channels=self.in_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
											]

											_SCREAMING_SNAKE_CASE					:			Dict							= []

											for _ in range(self.num_layers    ):
															_SCREAMING_SNAKE_CASE					:			str							= FlaxTransformeraDModel(
															    in_channels=self.in_channels       ,   n_heads=self.num_attention_heads       ,   d_head=self.in_channels // self.num_attention_heads       ,   depth=1       ,   use_linear_projection=self.use_linear_projection       ,   use_memory_efficient_attention=self.use_memory_efficient_attention       ,   dtype=self.dtype       ,   )
															attentions.append(__snake_case    )

															_SCREAMING_SNAKE_CASE					:			List[str]							= FlaxResnetBlockaD(
															    in_channels=self.in_channels       ,   out_channels=self.in_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
															resnets.append(__snake_case    )

											_SCREAMING_SNAKE_CASE					:			int							= resnets
											_SCREAMING_SNAKE_CASE					:			Dict							= attentions






							def __call__(       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case=True    ):
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= self.resnets[0](__snake_case       ,   __snake_case    )
											for attn, resnet in zip(self.attentions       ,   self.resnets[1:]    ):
															_SCREAMING_SNAKE_CASE					:			Tuple							= attn(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )
															_SCREAMING_SNAKE_CASE					:			Tuple							= resnet(__snake_case       ,   __snake_case       ,   deterministic=__snake_case    )

											return hidden_states








'''simple docstring'''


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


UpperCAmelCase_						:  List[Any]				  =						logging.get_logger(__name__)

UpperCAmelCase_						:  Any				  =						{
    'bigcode/gpt_bigcode-santacoder': 'https://huggingface.co/bigcode/gpt_bigcode-santacoder/resolve/main/config.json',
}





class        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: Dict       =    """gpt_bigcode"""
							A_				: str       =    ["""past_key_values"""]
							A_				: Tuple       =    {
							    """hidden_size""": """n_embd""",
							    """max_position_embeddings""": """n_positions""",
							    """num_attention_heads""": """n_head""",
							    """num_hidden_layers""": """n_layer""",
							}






							def __init__(       self       ,   __snake_case=5_0257       ,   __snake_case=1024       ,   __snake_case=768       ,   __snake_case=12       ,   __snake_case=12       ,   __snake_case=None       ,   __snake_case="gelu_pytorch_tanh"       ,   __snake_case=0.1       ,   __snake_case=0.1       ,   __snake_case=0.1       ,   __snake_case=1e-5       ,   __snake_case=0.02       ,   __snake_case=True       ,   __snake_case=True       ,   __snake_case=5_0256       ,   __snake_case=5_0256       ,   __snake_case=True       ,   __snake_case=True       ,   __snake_case=True       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= vocab_size
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= n_positions
											_SCREAMING_SNAKE_CASE					:			List[Any]							= n_embd
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= n_layer
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= n_head
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= n_inner
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= activation_function
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= resid_pdrop
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= embd_pdrop
											_SCREAMING_SNAKE_CASE					:			List[Any]							= attn_pdrop
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= layer_norm_epsilon
											_SCREAMING_SNAKE_CASE					:			Any							= initializer_range
											_SCREAMING_SNAKE_CASE					:			str							= scale_attn_weights
											_SCREAMING_SNAKE_CASE					:			Tuple							= use_cache
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= attention_softmax_in_fpaa
											_SCREAMING_SNAKE_CASE					:			List[Any]							= scale_attention_softmax_in_fpaa
											_SCREAMING_SNAKE_CASE					:			Any							= multi_query

											_SCREAMING_SNAKE_CASE					:			Optional[int]							= bos_token_id
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= eos_token_id

											super().__init__(bos_token_id=__snake_case       ,   eos_token_id=__snake_case       ,   **__snake_case    )




