code: 
  
  
  '''simple docstring'''
  
  
  
  
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  
  
  
  
  def 	UpperCAmelCase__(							_SCREAMING_SNAKE_CASE				:   str,_SCREAMING_SNAKE_CASE				:   int,_SCREAMING_SNAKE_CASE				:   Optional[int],_SCREAMING_SNAKE_CASE				:   Union[str, Any],_SCREAMING_SNAKE_CASE				:   int = 4096   ):
  
  
  
  
  
  
  			"""simple docstring"""
  
  
  
  
  			__A,  __A,  __A=  key.shape[-3:]
  			__A=  value.shape[-1]
  			__A=  min(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE   )
  			__A=  query / jnp.sqrt(_SCREAMING_SNAKE_CASE   )
  
  			@functools.partial(jax.checkpoint,prevent_cse=_SCREAMING_SNAKE_CASE   )
  			def summarize_chunk(_SCREAMING_SNAKE_CASE				:   Any,_SCREAMING_SNAKE_CASE				:   Dict,_SCREAMING_SNAKE_CASE				:   List[str]   ):
  						__A=  jnp.einsum('...qhd,...khd->...qhk',_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,precision=_SCREAMING_SNAKE_CASE   )
  
  						__A=  jnp.max(_SCREAMING_SNAKE_CASE,axis=-1,keepdims=_SCREAMING_SNAKE_CASE   )
  						__A=  jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE   )
  						__A=  jnp.exp(attn_weights - max_score   )
  
  						__A=  jnp.einsum('...vhf,...qhv->...qhf',_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,precision=_SCREAMING_SNAKE_CASE   )
  						__A=  jnp.einsum('...qhk->...qh',_SCREAMING_SNAKE_CASE   )
  
  						return (exp_values, exp_weights.sum(axis=-1   ), max_score)
  
  			def chunk_scanner(_SCREAMING_SNAKE_CASE				:   int   ):
  						# julienne key array
  						__A=  jax.lax.dynamic_slice(
  						    operand=_SCREAMING_SNAKE_CASE,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=_SCREAMING_SNAKE_CASE,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(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE   )
  
  			__A,  __A,  __A=  jax.lax.map(f=_SCREAMING_SNAKE_CASE,xs=jnp.arange(0,_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE   )   )
  
  			__A=  jnp.max(_SCREAMING_SNAKE_CASE,axis=0,keepdims=_SCREAMING_SNAKE_CASE   )
  			__A=  jnp.exp(chunk_max - global_max   )
  
  			chunk_values *= jnp.expand_dims(_SCREAMING_SNAKE_CASE,axis=-1   )
  			chunk_weights *= max_diffs
  
  			__A=  chunk_values.sum(axis=0   )
  			__A=  jnp.expand_dims(_SCREAMING_SNAKE_CASE,-1   ).sum(axis=0   )
  
  			return all_values / all_weights
  
  
  
  
  
  
  def 	UpperCAmelCase__(							_SCREAMING_SNAKE_CASE				:   List[str],_SCREAMING_SNAKE_CASE				:   Optional[Any],_SCREAMING_SNAKE_CASE				:   str,_SCREAMING_SNAKE_CASE				:   Optional[int]=jax.lax.Precision.HIGHEST,_SCREAMING_SNAKE_CASE				:   int = 1024,_SCREAMING_SNAKE_CASE				:   int = 4096   ):
  
  
  
  
  
  
  			"""simple docstring"""
  
  
  
  
  			__A,  __A,  __A=  query.shape[-3:]
  
  			def chunk_scanner(_SCREAMING_SNAKE_CASE				:   Any,_SCREAMING_SNAKE_CASE				:   int   ):
  						# julienne query array
  						__A=  jax.lax.dynamic_slice(
  						    operand=_SCREAMING_SNAKE_CASE,start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0],slice_sizes=list(query.shape[:-3]   ) + [min(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE   ), num_heads, q_features],)
  
  						return (
  						    chunk_idx + query_chunk_size,  # unused ignore it
  						    _query_chunk_attention(
  						        query=_SCREAMING_SNAKE_CASE,key=_SCREAMING_SNAKE_CASE,value=_SCREAMING_SNAKE_CASE,precision=_SCREAMING_SNAKE_CASE,key_chunk_size=_SCREAMING_SNAKE_CASE   ),
  						)
  
  			__A,  __A=  jax.lax.scan(
  			    f=_SCREAMING_SNAKE_CASE,init=0,xs=_SCREAMING_SNAKE_CASE,length=math.ceil(num_q / query_chunk_size   )   )
  
  			return jnp.concatenate(_SCREAMING_SNAKE_CASE,axis=-3   )  # fuse the chunked result back
  class 	a__				(				nn.Module      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							int
  			A	:							int							= 8
  			A	:							int							= 64
  			A	:							float							= 0.0
  			A	:							bool							= False
  			A	:							jnp.dtype							= jnp.floataa
  			def 				lowerCAmelCase    (     self       : Tuple  )			->		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(lowerCAmelCase_ ,						use_bias=lowerCAmelCase_ ,						dtype=self.dtype ,						name='to_q'  )
  						__A=  nn.Dense(lowerCAmelCase_ ,						use_bias=lowerCAmelCase_ ,						dtype=self.dtype ,						name='to_k'  )
  						__A=  nn.Dense(lowerCAmelCase_ ,						use_bias=lowerCAmelCase_ ,						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 				lowerCAmelCase    (     self       : Any ,						lowerCAmelCase_       : Optional[int]  )			->		str:
  						__A,  __A,  __A=  tensor.shape
  						__A=  self.heads
  						__A=  tensor.reshape(lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_ ,						dim // head_size  )
  						__A=  jnp.transpose(lowerCAmelCase_ ,						(0, 2, 1, 3)  )
  						__A=  tensor.reshape(batch_size * head_size ,						lowerCAmelCase_ ,						dim // head_size  )
  						return tensor
  			def 				lowerCAmelCase    (     self       : Tuple ,						lowerCAmelCase_       : Union[str, Any]  )			->		Any:
  						__A,  __A,  __A=  tensor.shape
  						__A=  self.heads
  						__A=  tensor.reshape(batch_size // head_size ,						lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_  )
  						__A=  jnp.transpose(lowerCAmelCase_ ,						(0, 2, 1, 3)  )
  						__A=  tensor.reshape(batch_size // head_size ,						lowerCAmelCase_ ,						dim * head_size  )
  						return tensor
  
  			def __call__(     self       : Dict ,						lowerCAmelCase_       : Union[str, Any] ,						lowerCAmelCase_       : Union[str, Any]=None ,						lowerCAmelCase_       : List[str]=True  )			->		int:
  						__A=  hidden_states if context is None else context
  
  						__A=  self.query(lowerCAmelCase_  )
  						__A=  self.key(lowerCAmelCase_  )
  						__A=  self.value(lowerCAmelCase_  )
  
  						__A=  self.reshape_heads_to_batch_dim(lowerCAmelCase_  )
  						__A=  self.reshape_heads_to_batch_dim(lowerCAmelCase_  )
  						__A=  self.reshape_heads_to_batch_dim(lowerCAmelCase_  )
  
  						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(lowerCAmelCase_  )
  
  									__A=  jax_memory_efficient_attention(
  									    lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_ ,						query_chunk_size=lowerCAmelCase_ ,						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' ,						lowerCAmelCase_ ,						lowerCAmelCase_  )
  									__A=  attention_scores * self.scale
  									__A=  nn.softmax(lowerCAmelCase_ ,						axis=2  )
  
  									# attend to values
  									__A=  jnp.einsum('b i j, b j d -> b i d' ,						lowerCAmelCase_ ,						lowerCAmelCase_  )
  
  						__A=  self.reshape_batch_dim_to_heads(lowerCAmelCase_  )
  						__A=  self.proj_attn(lowerCAmelCase_  )
  						return self.dropout_layer(lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  
  
  class 	a__				(				nn.Module      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							int
  			A	:							int
  			A	:							int
  			A	:							float							= 0.0
  			A	:							bool							= False
  			A	:							jnp.dtype							= jnp.floataa
  			A	:							bool							= False
  			def 				lowerCAmelCase    (     self       : Optional[int]  )			->		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       : int ,						lowerCAmelCase_       : List[Any] ,						lowerCAmelCase_       : List[str] ,						lowerCAmelCase_       : Any=True  )			->		List[str]:
  						# self attention
  						__A=  hidden_states
  						if self.only_cross_attention:
  									__A=  self.attna(self.norma(lowerCAmelCase_  ) ,						lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  						else:
  									__A=  self.attna(self.norma(lowerCAmelCase_  ) ,						deterministic=lowerCAmelCase_  )
  						__A=  hidden_states + residual
  
  						# cross attention
  						__A=  hidden_states
  						__A=  self.attna(self.norma(lowerCAmelCase_  ) ,						lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  						__A=  hidden_states + residual
  
  						# feed forward
  						__A=  hidden_states
  						__A=  self.ff(self.norma(lowerCAmelCase_  ) ,						deterministic=lowerCAmelCase_  )
  						__A=  hidden_states + residual
  
  						return self.dropout_layer(lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  
  
  class 	a__				(				nn.Module      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							int
  			A	:							int
  			A	:							int
  			A	:							int							= 1
  			A	:							float							= 0.0
  			A	:							bool							= False
  			A	:							bool							= False
  			A	:							jnp.dtype							= jnp.floataa
  			A	:							bool							= False
  			def 				lowerCAmelCase    (     self       : List[Any]  )			->		Any:
  						__A=  nn.GroupNorm(num_groups=32 ,						epsilon=1E-5  )
  
  						__A=  self.n_heads * self.d_head
  						if self.use_linear_projection:
  									__A=  nn.Dense(lowerCAmelCase_ ,						dtype=self.dtype  )
  						else:
  									__A=  nn.Conv(
  									    lowerCAmelCase_ ,						kernel_size=(1, 1) ,						strides=(1, 1) ,						padding='VALID' ,						dtype=self.dtype ,						)
  
  						__A=  [
  						    FlaxBasicTransformerBlock(
  						        lowerCAmelCase_ ,						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(lowerCAmelCase_ ,						dtype=self.dtype  )
  						else:
  									__A=  nn.Conv(
  									    lowerCAmelCase_ ,						kernel_size=(1, 1) ,						strides=(1, 1) ,						padding='VALID' ,						dtype=self.dtype ,						)
  
  						__A=  nn.Dropout(rate=self.dropout  )
  
  			def __call__(     self       : Optional[Any] ,						lowerCAmelCase_       : Union[str, Any] ,						lowerCAmelCase_       : Any ,						lowerCAmelCase_       : int=True  )			->		Tuple:
  						__A,  __A,  __A,  __A=  hidden_states.shape
  						__A=  hidden_states
  						__A=  self.norm(lowerCAmelCase_  )
  						if self.use_linear_projection:
  									__A=  hidden_states.reshape(lowerCAmelCase_ ,						height * width ,						lowerCAmelCase_  )
  									__A=  self.proj_in(lowerCAmelCase_  )
  						else:
  									__A=  self.proj_in(lowerCAmelCase_  )
  									__A=  hidden_states.reshape(lowerCAmelCase_ ,						height * width ,						lowerCAmelCase_  )
  
  						for transformer_block in self.transformer_blocks:
  									__A=  transformer_block(lowerCAmelCase_ ,						lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  
  						if self.use_linear_projection:
  									__A=  self.proj_out(lowerCAmelCase_  )
  									__A=  hidden_states.reshape(lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_  )
  						else:
  									__A=  hidden_states.reshape(lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_ ,						lowerCAmelCase_  )
  									__A=  self.proj_out(lowerCAmelCase_  )
  
  						__A=  hidden_states + residual
  						return self.dropout_layer(lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  
  
  class 	a__				(				nn.Module      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							int
  			A	:							float							= 0.0
  			A	:							jnp.dtype							= jnp.floataa
  			def 				lowerCAmelCase    (     self       : List[Any]  )			->		Any:
  						# 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       : List[Any] ,						lowerCAmelCase_       : Union[str, Any] ,						lowerCAmelCase_       : Union[str, Any]=True  )			->		Union[str, Any]:
  						__A=  self.net_a(lowerCAmelCase_ ,						deterministic=lowerCAmelCase_  )
  						__A=  self.net_a(lowerCAmelCase_  )
  						return hidden_states
  
  
  class 	a__				(				nn.Module      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							int
  			A	:							float							= 0.0
  			A	:							jnp.dtype							= jnp.floataa
  			def 				lowerCAmelCase    (     self       : Tuple  )			->		str:
  						__A=  self.dim * 4
  						__A=  nn.Dense(inner_dim * 2 ,						dtype=self.dtype  )
  						__A=  nn.Dropout(rate=self.dropout  )
  
  			def __call__(     self       : List[str] ,						lowerCAmelCase_       : Tuple ,						lowerCAmelCase_       : List[Any]=True  )			->		Any:
  						__A=  self.proj(lowerCAmelCase_  )
  						__A,  __A=  jnp.split(lowerCAmelCase_ ,						2 ,						axis=2  )
  						return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase_  ) ,						deterministic=lowerCAmelCase_  )
  
  
  
  
style_context: 
  
  
  '''simple docstring'''
  
  
  
  
  
  from ...configuration_utils import PretrainedConfig
  from ...utils import logging
  
  
  UpperCAmelCase__      =							logging.get_logger(__name__)
  
  UpperCAmelCase__      =							{
      '''dandelin/vilt-b32-mlm''': '''https://huggingface.co/dandelin/vilt-b32-mlm/blob/main/config.json'''
  }
  class 	a__				(				a_      ):
  
  
  
  
  
  			'''simple docstring'''
  
  			A	:							Dict							= '''vilt'''
  			def __init__(     self       : str ,						lowerCAmelCase_       : Tuple=30_522 ,						lowerCAmelCase_       : List[Any]=2 ,						lowerCAmelCase_       : Optional[Any]=2 ,						lowerCAmelCase_       : Union[str, Any]=40 ,						lowerCAmelCase_       : str=768 ,						lowerCAmelCase_       : Tuple=12 ,						lowerCAmelCase_       : List[Any]=12 ,						lowerCAmelCase_       : int=3_072 ,						lowerCAmelCase_       : List[Any]="gelu" ,						lowerCAmelCase_       : int=0.0 ,						lowerCAmelCase_       : str=0.0 ,						lowerCAmelCase_       : Optional[Any]=0.02 ,						lowerCAmelCase_       : Union[str, Any]=1E-12 ,						lowerCAmelCase_       : Dict=384 ,						lowerCAmelCase_       : Optional[int]=32 ,						lowerCAmelCase_       : Optional[Any]=3 ,						lowerCAmelCase_       : str=True ,						lowerCAmelCase_       : List[Any]=-1 ,						lowerCAmelCase_       : Union[str, Any]=False ,						lowerCAmelCase_       : str=-1 ,						**lowerCAmelCase_       : Any ,						)			->		Tuple:
  						super().__init__(tie_word_embeddings=lowerCAmelCase_ ,						**lowerCAmelCase_  )
  
  						__A=  vocab_size
  						__A=  type_vocab_size
  						__A=  modality_type_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=  layer_norm_eps
  
  						__A=  image_size
  						__A=  patch_size
  						__A=  num_channels
  						__A=  qkv_bias
  						__A=  max_image_length
  						__A=  num_images
  
  
  
  