code: 
  
  
  """simple docstring"""
  
  
  
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  
  def 						__a      (  a,		a,		a,		a,		a = 4_0_9_6					):
  
  
  
  
  
         """simple docstring"""
  
  
  
  
  
  
  
         _a   , _a   , _a          =				key.shape[-3:]
         _a          =				value.shape[-1]
         _a          =				min(a,		a					)
         _a          =				query / jnp.sqrt(a					)
  
         @functools.partial(jax.checkpoint,		prevent_cse=a					)
         def summarize_chunk(a,		a,		a					):
                _a          =				jnp.einsum("...qhd,...khd->...qhk",		a,		a,		precision=a					)
  
                _a          =				jnp.max(a,		axis=-1,		keepdims=a					)
                _a          =				jax.lax.stop_gradient(a					)
                _a          =				jnp.exp(attn_weights - max_score					)
  
                _a          =				jnp.einsum("...vhf,...qhv->...qhf",		a,		a,		precision=a					)
                _a          =				jnp.einsum("...qhk->...qh",		a					)
  
                return (exp_values, exp_weights.sum(axis=-1					), max_score)
  
         def chunk_scanner(a					):
                # julienne key array
                _a          =				jax.lax.dynamic_slice(
                    operand=a,		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
                _a          =				jax.lax.dynamic_slice(
                    operand=a,		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(a,		a,		a					)
  
         _a   , _a   , _a          =				jax.lax.map(f=a,		xs=jnp.arange(0,		a,		a					)					)
  
         _a          =				jnp.max(a,		axis=0,		keepdims=a					)
         _a          =				jnp.exp(chunk_max - global_max					)
  
         chunk_values *= jnp.expand_dims(a,		axis=-1					)
         chunk_weights *= max_diffs
  
         _a          =				chunk_values.sum(axis=0					)
         _a          =				jnp.expand_dims(a,		-1					).sum(axis=0					)
  
         return all_values / all_weights
  
  
  
  def 						__a      (  a,		a,		a,		a=jax.lax.Precision.HIGHEST,		a = 1_0_2_4,		a = 4_0_9_6					):
  
  
  
  
  
         """simple docstring"""
  
  
  
  
  
  
  
         _a   , _a   , _a          =				query.shape[-3:]
  
         def chunk_scanner(a,		a					):
                # julienne query array
                _a          =				jax.lax.dynamic_slice(
                    operand=a,		start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0],		slice_sizes=list(query.shape[:-3]					) + [min(a,		a					), num_heads, q_features],		)
  
                return (
                    chunk_idx + query_chunk_size,  # unused ignore it
                    _query_chunk_attention(
                        query=a,		key=a,		value=a,		precision=a,		key_chunk_size=a					),
                )
  
         _a   , _a          =				jax.lax.scan(
             f=a,		init=0,		xs=a,		length=math.ceil(num_q / query_chunk_size					)					)
  
         return jnp.concatenate(a,		axis=-3					)  # fuse the chunked result back
  
  
  
  
  
  class __snake_case		(					nn.Module      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int     	=			8
     lowerCAmelCase_				:   int     	=			64
     lowerCAmelCase_				:   float     	=			0.0
     lowerCAmelCase_				:   bool     	=			False
     lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:Any       ):
            _a          =				self.dim_head * self.heads
            _a          =				self.dim_head**-0.5
  
            # Weights were exported with old names {to_q, to_k, to_v, to_out}
            _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_q"       )
            _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_k"       )
            _a          =				nn.Dense(UpperCamelCase__				,  use_bias=UpperCamelCase__				,  dtype=self.dtype				,  name="to_v"       )
  
            _a          =				nn.Dense(self.query_dim				,  dtype=self.dtype				,  name="to_out_0"       )
            _a          =				nn.Dropout(rate=self.dropout       )
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:List[str]				,  UpperCamelCase__							:Tuple       ):
            _a   , _a   , _a          =				tensor.shape
            _a          =				self.heads
            _a          =				tensor.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  dim // head_size       )
            _a          =				jnp.transpose(UpperCamelCase__				,  (0, 2, 1, 3)       )
            _a          =				tensor.reshape(batch_size * head_size				,  UpperCamelCase__				,  dim // head_size       )
            return tensor
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:int				,  UpperCamelCase__							:int       ):
            _a   , _a   , _a          =				tensor.shape
            _a          =				self.heads
            _a          =				tensor.reshape(batch_size // head_size				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
            _a          =				jnp.transpose(UpperCamelCase__				,  (0, 2, 1, 3)       )
            _a          =				tensor.reshape(batch_size // head_size				,  UpperCamelCase__				,  dim * head_size       )
            return tensor
  
  
  
  
     def __call__(		self							:Optional[int]				,  UpperCamelCase__							:Optional[int]				,  UpperCamelCase__							:int=None				,  UpperCamelCase__							:Optional[Any]=True       ):
            _a          =				hidden_states if context is None else context
  
            _a          =				self.query(UpperCamelCase__       )
            _a          =				self.key(UpperCamelCase__       )
            _a          =				self.value(UpperCamelCase__       )
  
            _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )
            _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )
            _a          =				self.reshape_heads_to_batch_dim(UpperCamelCase__       )
  
            if self.use_memory_efficient_attention:
                   _a          =				query_states.transpose(1				,  0				,  2       )
                   _a          =				key_states.transpose(1				,  0				,  2       )
                   _a          =				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
  
                   _a          =				query_states.shape[-3]
                   if flatten_latent_dim % 64 == 0:
                          _a          =				int(flatten_latent_dim / 64       )
                   elif flatten_latent_dim % 16 == 0:
                          _a          =				int(flatten_latent_dim / 16       )
                   elif flatten_latent_dim % 4 == 0:
                          _a          =				int(flatten_latent_dim / 4       )
                   else:
                          _a          =				int(UpperCamelCase__       )
  
                   _a          =				jax_memory_efficient_attention(
                       UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  query_chunk_size=UpperCamelCase__				,  key_chunk_size=4_096 * 4       )
  
                   _a          =				hidden_states.transpose(1				,  0				,  2       )
            else:
                   # compute attentions
                   _a          =				jnp.einsum("b i d, b j d->b i j"				,  UpperCamelCase__				,  UpperCamelCase__       )
                   _a          =				attention_scores * self.scale
                   _a          =				nn.softmax(UpperCamelCase__				,  axis=2       )
  
                   # attend to values
                   _a          =				jnp.einsum("b i j, b j d -> b i d"				,  UpperCamelCase__				,  UpperCamelCase__       )
  
            _a          =				self.reshape_batch_dim_to_heads(UpperCamelCase__       )
            _a          =				self.proj_attn(UpperCamelCase__       )
            return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )
  
  
  
  
  
  
  class __snake_case		(					nn.Module      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   float     	=			0.0
     lowerCAmelCase_				:   bool     	=			False
     lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
     lowerCAmelCase_				:   bool     	=			False
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:List[str]       ):
            # self attention (or cross_attention if only_cross_attention is True)
            _a          =				FlaxAttention(
                self.dim				,  self.n_heads				,  self.d_head				,  self.dropout				,  self.use_memory_efficient_attention				,  dtype=self.dtype       )
            # cross attention
            _a          =				FlaxAttention(
                self.dim				,  self.n_heads				,  self.d_head				,  self.dropout				,  self.use_memory_efficient_attention				,  dtype=self.dtype       )
            _a          =				FlaxFeedForward(dim=self.dim				,  dropout=self.dropout				,  dtype=self.dtype       )
            _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
            _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
            _a          =				nn.LayerNorm(epsilon=1E-5				,  dtype=self.dtype       )
            _a          =				nn.Dropout(rate=self.dropout       )
  
  
  
  
     def __call__(		self							:Union[str, Any]				,  UpperCamelCase__							:int				,  UpperCamelCase__							:str				,  UpperCamelCase__							:List[Any]=True       ):
            # self attention
            _a          =				hidden_states
            if self.only_cross_attention:
                   _a          =				self.attna(self.norma(UpperCamelCase__       )				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )
            else:
                   _a          =				self.attna(self.norma(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )
            _a          =				hidden_states + residual
  
            # cross attention
            _a          =				hidden_states
            _a          =				self.attna(self.norma(UpperCamelCase__       )				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )
            _a          =				hidden_states + residual
  
            # feed forward
            _a          =				hidden_states
            _a          =				self.ff(self.norma(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )
            _a          =				hidden_states + residual
  
            return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )
  
  
  
  
  
  
  class __snake_case		(					nn.Module      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   int     	=			1
     lowerCAmelCase_				:   float     	=			0.0
     lowerCAmelCase_				:   bool     	=			False
     lowerCAmelCase_				:   bool     	=			False
     lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
     lowerCAmelCase_				:   bool     	=			False
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:str       ):
            _a          =				nn.GroupNorm(num_groups=32				,  epsilon=1E-5       )
  
            _a          =				self.n_heads * self.d_head
            if self.use_linear_projection:
                   _a          =				nn.Dense(UpperCamelCase__				,  dtype=self.dtype       )
            else:
                   _a          =				nn.Conv(
                       UpperCamelCase__				,  kernel_size=(1, 1)				,  strides=(1, 1)				,  padding="VALID"				,  dtype=self.dtype				,  )
  
            _a          =				[
                FlaxBasicTransformerBlock(
                    UpperCamelCase__				,  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:
                   _a          =				nn.Dense(UpperCamelCase__				,  dtype=self.dtype       )
            else:
                   _a          =				nn.Conv(
                       UpperCamelCase__				,  kernel_size=(1, 1)				,  strides=(1, 1)				,  padding="VALID"				,  dtype=self.dtype				,  )
  
            _a          =				nn.Dropout(rate=self.dropout       )
  
  
  
  
     def __call__(		self							:Union[str, Any]				,  UpperCamelCase__							:str				,  UpperCamelCase__							:str				,  UpperCamelCase__							:Any=True       ):
            _a   , _a   , _a   , _a          =				hidden_states.shape
            _a          =				hidden_states
            _a          =				self.norm(UpperCamelCase__       )
            if self.use_linear_projection:
                   _a          =				hidden_states.reshape(UpperCamelCase__				,  height * width				,  UpperCamelCase__       )
                   _a          =				self.proj_in(UpperCamelCase__       )
            else:
                   _a          =				self.proj_in(UpperCamelCase__       )
                   _a          =				hidden_states.reshape(UpperCamelCase__				,  height * width				,  UpperCamelCase__       )
  
            for transformer_block in self.transformer_blocks:
                   _a          =				transformer_block(UpperCamelCase__				,  UpperCamelCase__				,  deterministic=UpperCamelCase__       )
  
            if self.use_linear_projection:
                   _a          =				self.proj_out(UpperCamelCase__       )
                   _a          =				hidden_states.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
            else:
                   _a          =				hidden_states.reshape(UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__				,  UpperCamelCase__       )
                   _a          =				self.proj_out(UpperCamelCase__       )
  
            _a          =				hidden_states + residual
            return self.dropout_layer(UpperCamelCase__				,  deterministic=UpperCamelCase__       )
  
  
  
  
  
  
  class __snake_case		(					nn.Module      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   float     	=			0.0
     lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:Tuple       ):
            # The second linear layer needs to be called
            # net_2 for now to match the index of the Sequential layer
            _a          =				FlaxGEGLU(self.dim				,  self.dropout				,  self.dtype       )
            _a          =				nn.Dense(self.dim				,  dtype=self.dtype       )
  
  
  
  
     def __call__(		self							:Any				,  UpperCamelCase__							:Optional[int]				,  UpperCamelCase__							:Union[str, Any]=True       ):
            _a          =				self.net_a(UpperCamelCase__				,  deterministic=UpperCamelCase__       )
            _a          =				self.net_a(UpperCamelCase__       )
            return hidden_states
  
  
  
  
  
  
  class __snake_case		(					nn.Module      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   int
     lowerCAmelCase_				:   float     	=			0.0
     lowerCAmelCase_				:   jnp.dtype     	=			jnp.floataa
  
  
  
  
     def SCREAMING_SNAKE_CASE_      (		self							:Union[str, Any]       ):
            _a          =				self.dim * 4
            _a          =				nn.Dense(inner_dim * 2				,  dtype=self.dtype       )
            _a          =				nn.Dropout(rate=self.dropout       )
  
  
  
  
     def __call__(		self							:Optional[int]				,  UpperCamelCase__							:Dict				,  UpperCamelCase__							:Union[str, Any]=True       ):
            _a          =				self.proj(UpperCamelCase__       )
            _a   , _a          =				jnp.split(UpperCamelCase__				,  2				,  axis=2       )
            return self.dropout_layer(hidden_linear * nn.gelu(UpperCamelCase__       )				,  deterministic=UpperCamelCase__       )
  
style_context: 
  
  
  """simple docstring"""
  
  
  
  
  from ...configuration_utils import PretrainedConfig
  from ...utils import logging
  
  
  __SCREAMING_SNAKE_CASE					=						logging.get_logger(__name__)
  
  __SCREAMING_SNAKE_CASE					=						{
      """bert-base-uncased""": """https://huggingface.co/bert-base-uncased/resolve/main/config.json""",
      # QDQBERT models can be loaded from any BERT checkpoint, available at https://huggingface.co/models?filter=bert
  }
  
  
  
  
  
  class __snake_case		(					_SCREAMING_SNAKE_CASE      ):
  
  
  
  
  
  
  
     """simple docstring"""
  
  
  
  
     lowerCAmelCase_				:   Tuple     	=			'qdqbert'
  
  
  
  
     def __init__(		self							:List[Any]				,  UpperCamelCase__							:List[Any]=30_522				,  UpperCamelCase__							:Optional[int]=768				,  UpperCamelCase__							:str=12				,  UpperCamelCase__							:Tuple=12				,  UpperCamelCase__							:Any=3_072				,  UpperCamelCase__							:Union[str, Any]="gelu"				,  UpperCamelCase__							:Dict=0.1				,  UpperCamelCase__							:Any=0.1				,  UpperCamelCase__							:Dict=512				,  UpperCamelCase__							:Optional[int]=2				,  UpperCamelCase__							:Any=0.02				,  UpperCamelCase__							:Optional[int]=1E-12				,  UpperCamelCase__							:str=True				,  UpperCamelCase__							:str=1				,  UpperCamelCase__							:List[Any]=0				,  UpperCamelCase__							:List[str]=2				,  **UpperCamelCase__							:Optional[Any]				,  ):
            super().__init__(pad_token_id=UpperCamelCase__				,  bos_token_id=UpperCamelCase__				,  eos_token_id=UpperCamelCase__				,  **UpperCamelCase__       )
  
            _a          =				vocab_size
            _a          =				max_position_embeddings
            _a          =				hidden_size
            _a          =				num_hidden_layers
            _a          =				num_attention_heads
            _a          =				intermediate_size
            _a          =				hidden_act
            _a          =				hidden_dropout_prob
            _a          =				attention_probs_dropout_prob
            _a          =				initializer_range
            _a          =				type_vocab_size
            _a          =				layer_norm_eps
            _a          =				use_cache
  