





from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


snake_case__					:							Optional[Any]													=     logging.get_logger(__name__)


snake_case__					:							Optional[Any]													=     {
    'Salesforce/codegen-350M-nl': 'https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json',
    'Salesforce/codegen-350M-multi': 'https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json',
    'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json',
    'Salesforce/codegen-2B-nl': 'https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json',
    'Salesforce/codegen-2B-multi': 'https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json',
    'Salesforce/codegen-2B-mono': 'https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json',
    'Salesforce/codegen-6B-nl': 'https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json',
    'Salesforce/codegen-6B-multi': 'https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json',
    'Salesforce/codegen-6B-mono': 'https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json',
    'Salesforce/codegen-16B-nl': 'https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json',
    'Salesforce/codegen-16B-multi': 'https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json',
    'Salesforce/codegen-16B-mono': 'https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json',
}





class A_					(			_lowerCamelCase   ):
   lowerCAmelCase__        = """codegen"""
   lowerCAmelCase__        = {
       """max_position_embeddings""": """n_positions""",
       """hidden_size""": """n_embd""",
       """num_attention_heads""": """n_head""",
       """num_hidden_layers""": """n_layer""",
   }




   def __init__(self	:int    ,   _UpperCamelCase	:Optional[int]=5_0400    ,   _UpperCamelCase	:Any=2048    ,   _UpperCamelCase	:Tuple=2048    ,   _UpperCamelCase	:Tuple=4096    ,   _UpperCamelCase	:Dict=28    ,   _UpperCamelCase	:Tuple=16    ,   _UpperCamelCase	:int=64    ,   _UpperCamelCase	:Optional[int]=None    ,   _UpperCamelCase	:int="gelu_new"    ,   _UpperCamelCase	:str=0.0    ,   _UpperCamelCase	:Dict=0.0    ,   _UpperCamelCase	:Dict=0.0    ,   _UpperCamelCase	:List[Any]=1e-5    ,   _UpperCamelCase	:Dict=0.0_2    ,   _UpperCamelCase	:Optional[Any]=True    ,   _UpperCamelCase	:Tuple=5_0256    ,   _UpperCamelCase	:str=5_0256    ,   _UpperCamelCase	:List[str]=False    ,   **_UpperCamelCase	:List[Any]    ,   )->					Union[str, Any]:
         __A          =				vocab_size
         __A          =				n_ctx
         __A          =				n_positions
         __A          =				n_embd
         __A          =				n_layer
         __A          =				n_head
         __A          =				n_inner
         __A          =				rotary_dim
         __A          =				activation_function
         __A          =				resid_pdrop
         __A          =				embd_pdrop
         __A          =				attn_pdrop
         __A          =				layer_norm_epsilon
         __A          =				initializer_range
         __A          =				use_cache

         __A          =				bos_token_id
         __A          =				eos_token_id

         super().__init__(
             bos_token_id=_UpperCamelCase    ,   eos_token_id=_UpperCamelCase    ,   tie_word_embeddings=_UpperCamelCase    ,   **_UpperCamelCase )





class A_					(			_lowerCamelCase   ):




   def __init__(self	:str    ,   _UpperCamelCase	:PretrainedConfig    ,   _UpperCamelCase	:str = "default"    ,   _UpperCamelCase	:List[PatchingSpec] = None    ,   _UpperCamelCase	:bool = False    ,   )->					Dict:
         super().__init__(_UpperCamelCase    ,   task=_UpperCamelCase    ,   patching_specs=_UpperCamelCase    ,   use_past=_UpperCamelCase )
         if not getattr(self._config    ,   '''pad_token_id'''    ,   _UpperCamelCase ):
               # TODO: how to do that better?
               __A          =				0




   @property
   def _lowerCAmelCase      (self	:Any )->					Mapping[str, Mapping[int, str]]:
         __A          =				OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}} )
         if self.use_past:
               self.fill_with_past_key_values_(_UpperCamelCase    ,   direction='''inputs''' )
               __A          =				{0: '''batch''', 1: '''past_sequence + sequence'''}
         else:
               __A          =				{0: '''batch''', 1: '''sequence'''}

         return common_inputs




   @property
   def _lowerCAmelCase      (self	:List[str] )->					int:
         return self._config.n_layer




   @property
   def _lowerCAmelCase      (self	:Optional[Any] )->					int:
         return self._config.n_head




   def _lowerCAmelCase      (self	:str    ,   _UpperCamelCase	:PreTrainedTokenizer    ,   _UpperCamelCase	:int = -1    ,   _UpperCamelCase	:int = -1    ,   _UpperCamelCase	:bool = False    ,   _UpperCamelCase	:Optional[TensorType] = None    ,   )->					Mapping[str, Any]:
         __A          =				super(_UpperCamelCase    ,   self ).generate_dummy_inputs(
             _UpperCamelCase    ,   batch_size=_UpperCamelCase    ,   seq_length=_UpperCamelCase    ,   is_pair=_UpperCamelCase    ,   framework=_UpperCamelCase )

         # We need to order the input in the way they appears in the forward()
         __A          =				OrderedDict({'''input_ids''': common_inputs['''input_ids''']} )

         # Need to add the past_keys
         if self.use_past:
               if not is_torch_available():
                     raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.''' )
               else:
                     import torch

                     __A   ,						__A          =				common_inputs['''input_ids'''].shape
                     # Not using the same length for past_key_values
                     __A          =				seqlen + 2
                     __A          =				(
                         batch,
                         self.num_attention_heads,
                         past_key_values_length,
                         self._config.hidden_size // self.num_attention_heads,
                     )
                     __A          =				[
                         (torch.zeros(_UpperCamelCase ), torch.zeros(_UpperCamelCase )) for _ in range(self.num_layers )
                     ]

         __A          =				common_inputs['''attention_mask''']
         if self.use_past:
               __A          =				ordered_inputs['''attention_mask'''].dtype
               __A          =				torch.cat(
                   [ordered_inputs['''attention_mask'''], torch.ones(_UpperCamelCase    ,   _UpperCamelCase    ,   dtype=_UpperCamelCase )]    ,   dim=1 )

         return ordered_inputs






   @property
   def _lowerCAmelCase      (self	:Any )->					int:
         return 13












from typing import Optional, Tuple, Union

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

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)





@flax.struct.dataclass
class A_					(			_lowerCamelCase   ):
   lowerCAmelCase__        = 42





@flax_register_to_config
class A_					(			nn.Module ,      _lowerCamelCase ,      _lowerCamelCase   ):
   lowerCAmelCase__        = 32
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   lowerCAmelCase__        = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   lowerCAmelCase__        = False
   lowerCAmelCase__        = (320, 640, 1280, 1280)
   lowerCAmelCase__        = 2
   lowerCAmelCase__        = 8
   lowerCAmelCase__        = None
   lowerCAmelCase__        = 1280
   lowerCAmelCase__        = 0.0
   lowerCAmelCase__        = False
   lowerCAmelCase__        = jnp.floataa
   lowerCAmelCase__        = True
   lowerCAmelCase__        = 0
   lowerCAmelCase__        = False




   def _lowerCAmelCase      (self	:Tuple    ,   _UpperCamelCase	:jax.random.KeyArray )->					FrozenDict:
         # init input tensors
         __A          =				(1, self.in_channels, self.sample_size, self.sample_size)
         __A          =				jnp.zeros(_UpperCamelCase    ,   dtype=jnp.floataa )
         __A          =				jnp.ones((1,)    ,   dtype=jnp.intaa )
         __A          =				jnp.zeros((1, 1, self.cross_attention_dim)    ,   dtype=jnp.floataa )

         __A   ,						__A          =				jax.random.split(_UpperCamelCase )
         __A          =				{'''params''': params_rng, '''dropout''': dropout_rng}

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




   def _lowerCAmelCase      (self	:Tuple )->					Optional[int]:
         __A          =				self.block_out_channels
         __A          =				block_out_channels[0] * 4

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

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

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

         # time
         __A          =				FlaxTimesteps(
             block_out_channels[0]    ,   flip_sin_to_cos=self.flip_sin_to_cos    ,   freq_shift=self.config.freq_shift )
         __A          =				FlaxTimestepEmbedding(_UpperCamelCase    ,   dtype=self.dtype )

         __A          =				self.only_cross_attention
         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(only_cross_attention,) * len(self.down_block_types )

         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(num_attention_heads,) * len(self.down_block_types )

         # down
         __A          =				[]
         __A          =				block_out_channels[0]
         for i, down_block_type in enumerate(self.down_block_types ):
               __A          =				output_channel
               __A          =				block_out_channels[i]
               __A          =				i == len(_UpperCamelCase ) - 1

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

               down_blocks.append(_UpperCamelCase )
         __A          =				down_blocks

         # mid
         __A          =				FlaxUNetMidBlockaDCrossAttn(
             in_channels=block_out_channels[-1]    ,   dropout=self.dropout    ,   num_attention_heads=num_attention_heads[-1]    ,   use_linear_projection=self.use_linear_projection    ,   use_memory_efficient_attention=self.use_memory_efficient_attention    ,   dtype=self.dtype    ,   )

         # up
         __A          =				[]
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				reversed_block_out_channels[0]
         for i, up_block_type in enumerate(self.up_block_types ):
               __A          =				output_channel
               __A          =				reversed_block_out_channels[i]
               __A          =				reversed_block_out_channels[min(i + 1    ,   len(_UpperCamelCase ) - 1 )]

               __A          =				i == len(_UpperCamelCase ) - 1

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

               up_blocks.append(_UpperCamelCase )
               __A          =				output_channel
         __A          =				up_blocks

         # out
         __A          =				nn.GroupNorm(num_groups=32    ,   epsilon=1e-5 )
         __A          =				nn.Conv(
             self.out_channels    ,   kernel_size=(3, 3)    ,   strides=(1, 1)    ,   padding=((1, 1), (1, 1))    ,   dtype=self.dtype    ,   )






   def __call__(self	:Union[str, Any]    ,   _UpperCamelCase	:Any    ,   _UpperCamelCase	:List[str]    ,   _UpperCamelCase	:str    ,   _UpperCamelCase	:int=None    ,   _UpperCamelCase	:Optional[int]=None    ,   _UpperCamelCase	:bool = True    ,   _UpperCamelCase	:bool = False    ,   )->					Union[FlaxUNetaDConditionOutput, Tuple]:
         # 1. time
         if not isinstance(_UpperCamelCase    ,   jnp.ndarray ):
               __A          =				jnp.array([timesteps]    ,   dtype=jnp.intaa )
         elif isinstance(_UpperCamelCase    ,   jnp.ndarray ) and len(timesteps.shape ) == 0:
               __A          =				timesteps.astype(dtype=jnp.floataa )
               __A          =				jnp.expand_dims(_UpperCamelCase    ,   0 )

         __A          =				self.time_proj(_UpperCamelCase )
         __A          =				self.time_embedding(_UpperCamelCase )

         # 2. pre-process
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 2, 3, 1) )
         __A          =				self.conv_in(_UpperCamelCase )

         # 3. down
         __A          =				(sample,)
         for down_block in self.down_blocks:
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               else:
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               down_block_res_samples += res_samples

         if down_block_additional_residuals is not None:
               __A          =				()

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

               __A          =				new_down_block_res_samples

         # 4. mid
         __A          =				self.mid_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )

         if mid_block_additional_residual is not None:
               sample += mid_block_additional_residual

         # 5. up
         for up_block in self.up_blocks:
               __A          =				down_block_res_samples[-(self.layers_per_block + 1) :]
               __A          =				down_block_res_samples[: -(self.layers_per_block + 1)]
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A          =				up_block(
                         _UpperCamelCase    ,   temb=_UpperCamelCase    ,   encoder_hidden_states=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train    ,   )
               else:
                     __A          =				up_block(_UpperCamelCase    ,   temb=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train )

        # 6. post-process
         __A          =				self.conv_norm_out(_UpperCamelCase )
         __A          =				nn.silu(_UpperCamelCase )
         __A          =				self.conv_out(_UpperCamelCase )
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 3, 1, 2) )

         if not return_dict:
               return (sample,)

         return FlaxUNetaDConditionOutput(sample=_UpperCamelCase )





