




"""simple docstring"""







import functools
import math

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





def 							lowercase   (snake_case__				:	Optional[Any] ,    snake_case__				:	Tuple ,    snake_case__				:	Optional[Any] ,    snake_case__				:	str ,    snake_case__				:	Union[str, Any] = 4_096	)       ->					Optional[Any]:


				'''simple docstring'''






				lowerCAmelCase		=			key.shape[-3:]
				lowerCAmelCase		=			value.shape[-1]
				lowerCAmelCase		=			min(_UpperCamelCase ,    _UpperCamelCase	)
				lowerCAmelCase		=			query / jnp.sqrt(_UpperCamelCase	)

				@functools.partial(jax.checkpoint ,    prevent_cse=_UpperCamelCase	)
				def summarize_chunk(snake_case__				:	Tuple ,    snake_case__				:	Dict ,    snake_case__				:	Any	):
								lowerCAmelCase		=			jnp.einsum("""...qhd,...khd->...qhk""" ,    _UpperCamelCase ,    _UpperCamelCase ,    precision=_UpperCamelCase	)

								lowerCAmelCase		=			jnp.max(_UpperCamelCase ,    axis=-1 ,    keepdims=_UpperCamelCase	)
								lowerCAmelCase		=			jax.lax.stop_gradient(_UpperCamelCase	)
								lowerCAmelCase		=			jnp.exp(attn_weights - max_score	)

								lowerCAmelCase		=			jnp.einsum("""...vhf,...qhv->...qhf""" ,    _UpperCamelCase ,    _UpperCamelCase ,    precision=_UpperCamelCase	)
								lowerCAmelCase		=			jnp.einsum("""...qhk->...qh""" ,    _UpperCamelCase	)

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

				def chunk_scanner(snake_case__				:	int	):
								# julienne key array
								lowerCAmelCase		=			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		=			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		=			jax.lax.map(f=_UpperCamelCase ,    xs=jnp.arange(0 ,    _UpperCamelCase ,    _UpperCamelCase	)	)

				lowerCAmelCase		=			jnp.max(_UpperCamelCase ,    axis=0 ,    keepdims=_UpperCamelCase	)
				lowerCAmelCase		=			jnp.exp(chunk_max - global_max	)

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

				lowerCAmelCase		=			chunk_values.sum(axis=0	)
				lowerCAmelCase		=			jnp.expand_dims(_UpperCamelCase ,    -1	).sum(axis=0	)

				return all_values / all_weights





def 							lowercase   (snake_case__				:	Optional[Any] ,    snake_case__				:	str ,    snake_case__				:	Any ,    snake_case__				:	Dict=jax.lax.Precision.HIGHEST ,    snake_case__				:	Optional[Any] = 1_024 ,    snake_case__				:	Optional[Any] = 4_096	)       ->					Any:


				'''simple docstring'''






				lowerCAmelCase		=			query.shape[-3:]

				def chunk_scanner(snake_case__				:	Optional[int] ,    snake_case__				:	Union[str, Any]	):
								# julienne query array
								lowerCAmelCase		=			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		=			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  SCREAMING_SNAKE_CASE__   (   nn.Module   ):
						_a        =  42
						_a        =  8
						_a        =  64
						_a        =  0.0
						_a        =  False
						_a        =  jnp.floataa


						def 					__lowercase	(							self			:     Dict   ):
										lowerCAmelCase		=			self.dim_head * self.heads
										lowerCAmelCase		=			self.dim_head**-0.5

										# Weights were exported with old names {to_q, to_k, to_v, to_out}
										lowerCAmelCase		=			nn.Dense(lowerCAmelCase  ,  use_bias=lowerCAmelCase  ,  dtype=self.dtype  ,  name="""to_q"""   )
										lowerCAmelCase		=			nn.Dense(lowerCAmelCase  ,  use_bias=lowerCAmelCase  ,  dtype=self.dtype  ,  name="""to_k"""   )
										lowerCAmelCase		=			nn.Dense(lowerCAmelCase  ,  use_bias=lowerCAmelCase  ,  dtype=self.dtype  ,  name="""to_v"""   )

										lowerCAmelCase		=			nn.Dense(self.query_dim  ,  dtype=self.dtype  ,  name="""to_out_0"""   )
										lowerCAmelCase		=			nn.Dropout(rate=self.dropout   )


						def 					__lowercase	(							self			:     Optional[int]  ,  lowerCAmelCase			:     Tuple   ):
										lowerCAmelCase		=			tensor.shape
										lowerCAmelCase		=			self.heads
										lowerCAmelCase		=			tensor.reshape(lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  dim // head_size   )
										lowerCAmelCase		=			jnp.transpose(lowerCAmelCase  ,  (0, 2, 1, 3)   )
										lowerCAmelCase		=			tensor.reshape(batch_size * head_size  ,  lowerCAmelCase  ,  dim // head_size   )
										return tensor


						def 					__lowercase	(							self			:     List[str]  ,  lowerCAmelCase			:     List[str]   ):
										lowerCAmelCase		=			tensor.shape
										lowerCAmelCase		=			self.heads
										lowerCAmelCase		=			tensor.reshape(batch_size // head_size  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase   )
										lowerCAmelCase		=			jnp.transpose(lowerCAmelCase  ,  (0, 2, 1, 3)   )
										lowerCAmelCase		=			tensor.reshape(batch_size // head_size  ,  lowerCAmelCase  ,  dim * head_size   )
										return tensor




						def __call__(							self			:     Dict  ,  lowerCAmelCase			:     Dict  ,  lowerCAmelCase			:     List[Any]=None  ,  lowerCAmelCase			:     int=True   ):
										lowerCAmelCase		=			hidden_states if context is None else context

										lowerCAmelCase		=			self.query(lowerCAmelCase   )
										lowerCAmelCase		=			self.key(lowerCAmelCase   )
										lowerCAmelCase		=			self.value(lowerCAmelCase   )

										lowerCAmelCase		=			self.reshape_heads_to_batch_dim(lowerCAmelCase   )
										lowerCAmelCase		=			self.reshape_heads_to_batch_dim(lowerCAmelCase   )
										lowerCAmelCase		=			self.reshape_heads_to_batch_dim(lowerCAmelCase   )

										if self.use_memory_efficient_attention:
														lowerCAmelCase		=			query_states.transpose(1  ,  0  ,  2   )
														lowerCAmelCase		=			key_states.transpose(1  ,  0  ,  2   )
														lowerCAmelCase		=			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		=			query_states.shape[-3]
														if flatten_latent_dim % 64 == 0:
																		lowerCAmelCase		=			int(flatten_latent_dim / 64   )
														elif flatten_latent_dim % 16 == 0:
																		lowerCAmelCase		=			int(flatten_latent_dim / 16   )
														elif flatten_latent_dim % 4 == 0:
																		lowerCAmelCase		=			int(flatten_latent_dim / 4   )
														else:
																		lowerCAmelCase		=			int(lowerCAmelCase   )

														lowerCAmelCase		=			jax_memory_efficient_attention(
														    lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  query_chunk_size=lowerCAmelCase  ,  key_chunk_size=4096 * 4   )

														lowerCAmelCase		=			hidden_states.transpose(1  ,  0  ,  2   )
										else:
														# compute attentions
														lowerCAmelCase		=			jnp.einsum("""b i d, b j d->b i j"""  ,  lowerCAmelCase  ,  lowerCAmelCase   )
														lowerCAmelCase		=			attention_scores * self.scale
														lowerCAmelCase		=			nn.softmax(lowerCAmelCase  ,  axis=2   )

														# attend to values
														lowerCAmelCase		=			jnp.einsum("""b i j, b j d -> b i d"""  ,  lowerCAmelCase  ,  lowerCAmelCase   )

										lowerCAmelCase		=			self.reshape_batch_dim_to_heads(lowerCAmelCase   )
										lowerCAmelCase		=			self.proj_attn(lowerCAmelCase   )
										return self.dropout_layer(lowerCAmelCase  ,  deterministic=lowerCAmelCase   )






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


						def 					__lowercase	(							self			:     Any   ):
										lowerCAmelCase		=			FlaxAttention(
										    self.dim  ,  self.n_heads  ,  self.d_head  ,  self.dropout  ,  self.use_memory_efficient_attention  ,  dtype=self.dtype   )
										# cross attention
										lowerCAmelCase		=			FlaxAttention(
										    self.dim  ,  self.n_heads  ,  self.d_head  ,  self.dropout  ,  self.use_memory_efficient_attention  ,  dtype=self.dtype   )
										lowerCAmelCase		=			FlaxFeedForward(dim=self.dim  ,  dropout=self.dropout  ,  dtype=self.dtype   )
										lowerCAmelCase		=			nn.LayerNorm(epsilon=1e-5  ,  dtype=self.dtype   )
										lowerCAmelCase		=			nn.LayerNorm(epsilon=1e-5  ,  dtype=self.dtype   )
										lowerCAmelCase		=			nn.LayerNorm(epsilon=1e-5  ,  dtype=self.dtype   )
										lowerCAmelCase		=			nn.Dropout(rate=self.dropout   )




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

										# cross attention
										lowerCAmelCase		=			hidden_states
										lowerCAmelCase		=			self.attna(self.norma(lowerCAmelCase   )  ,  lowerCAmelCase  ,  deterministic=lowerCAmelCase   )
										lowerCAmelCase		=			hidden_states + residual

										# feed forward
										lowerCAmelCase		=			hidden_states
										lowerCAmelCase		=			self.ff(self.norma(lowerCAmelCase   )  ,  deterministic=lowerCAmelCase   )
										lowerCAmelCase		=			hidden_states + residual

										return self.dropout_layer(lowerCAmelCase  ,  deterministic=lowerCAmelCase   )






class  SCREAMING_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 					__lowercase	(							self			:     Optional[Any]   ):
										lowerCAmelCase		=			nn.GroupNorm(num_groups=32  ,  epsilon=1e-5   )

										lowerCAmelCase		=			self.n_heads * self.d_head
										if self.use_linear_projection:
														lowerCAmelCase		=			nn.Dense(lowerCAmelCase  ,  dtype=self.dtype   )
										else:
														lowerCAmelCase		=			nn.Conv(
														    lowerCAmelCase  ,  kernel_size=(1, 1)  ,  strides=(1, 1)  ,  padding="""VALID"""  ,  dtype=self.dtype  ,  )

										lowerCAmelCase		=			[
										    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:
														lowerCAmelCase		=			nn.Dense(lowerCAmelCase  ,  dtype=self.dtype   )
										else:
														lowerCAmelCase		=			nn.Conv(
														    lowerCAmelCase  ,  kernel_size=(1, 1)  ,  strides=(1, 1)  ,  padding="""VALID"""  ,  dtype=self.dtype  ,  )

										lowerCAmelCase		=			nn.Dropout(rate=self.dropout   )




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

										for transformer_block in self.transformer_blocks:
														lowerCAmelCase		=			transformer_block(lowerCAmelCase  ,  lowerCAmelCase  ,  deterministic=lowerCAmelCase   )

										if self.use_linear_projection:
														lowerCAmelCase		=			self.proj_out(lowerCAmelCase   )
														lowerCAmelCase		=			hidden_states.reshape(lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase   )
										else:
														lowerCAmelCase		=			hidden_states.reshape(lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase  ,  lowerCAmelCase   )
														lowerCAmelCase		=			self.proj_out(lowerCAmelCase   )

										lowerCAmelCase		=			hidden_states + residual
										return self.dropout_layer(lowerCAmelCase  ,  deterministic=lowerCAmelCase   )






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


						def 					__lowercase	(							self			:     Tuple   ):
										lowerCAmelCase		=			FlaxGEGLU(self.dim  ,  self.dropout  ,  self.dtype   )
										lowerCAmelCase		=			nn.Dense(self.dim  ,  dtype=self.dtype   )




						def __call__(							self			:     Tuple  ,  lowerCAmelCase			:     Optional[int]  ,  lowerCAmelCase			:     Any=True   ):
										lowerCAmelCase		=			self.net_a(lowerCAmelCase  ,  deterministic=lowerCAmelCase   )
										lowerCAmelCase		=			self.net_a(lowerCAmelCase   )
										return hidden_states






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


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




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






"""simple docstring"""







import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding






class  SCREAMING_SNAKE_CASE__   (   _a   ):
						_a        =  ['image_processor', 'tokenizer']
						_a        =  'VideoMAEImageProcessor'
						_a        =  ('CLIPTokenizer', 'CLIPTokenizerFast')


						def __init__(							self			:     Union[str, Any]  ,  lowerCAmelCase			:     Optional[Any]=None  ,  lowerCAmelCase			:     Tuple=None  ,  **lowerCAmelCase			:     List[str]   ):
										lowerCAmelCase		=			None
										if "feature_extractor" in kwargs:
														warnings.warn(
														    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
														    """ instead."""  ,  lowerCAmelCase  ,  )
														lowerCAmelCase		=			kwargs.pop("""feature_extractor"""   )

										lowerCAmelCase		=			image_processor if image_processor is not None else feature_extractor
										if image_processor is None:
														raise ValueError("""You need to specify an `image_processor`."""   )
										if tokenizer is None:
														raise ValueError("""You need to specify a `tokenizer`."""   )

										super().__init__(lowerCAmelCase  ,  lowerCAmelCase   )
										lowerCAmelCase		=			self.image_processor


						def __call__(							self			:     Union[str, Any]  ,  lowerCAmelCase			:     str=None  ,  lowerCAmelCase			:     Tuple=None  ,  lowerCAmelCase			:     int=None  ,  **lowerCAmelCase			:     str   ):

										if text is None and videos is None:
														raise ValueError("""You have to specify either text or videos. Both cannot be none."""   )

										if text is not None:
														lowerCAmelCase		=			self.tokenizer(lowerCAmelCase  ,  return_tensors=lowerCAmelCase  ,  **lowerCAmelCase   )

										if videos is not None:
														lowerCAmelCase		=			self.image_processor(lowerCAmelCase  ,  return_tensors=lowerCAmelCase  ,  **lowerCAmelCase   )

										if text is not None and videos is not None:
														lowerCAmelCase		=			image_features.pixel_values
														return encoding
										elif text is not None:
														return encoding
										else:
														return BatchEncoding(data=dict(**lowerCAmelCase   )  ,  tensor_type=lowerCAmelCase   )


						def 					__lowercase	(							self			:     str  ,  *lowerCAmelCase			:     List[str]  ,  **lowerCAmelCase			:     int   ):
										return self.tokenizer.batch_decode(*lowerCAmelCase  ,  **lowerCAmelCase   )


						def 					__lowercase	(							self			:     List[str]  ,  *lowerCAmelCase			:     Any  ,  **lowerCAmelCase			:     List[Any]   ):
										return self.tokenizer.decode(*lowerCAmelCase  ,  **lowerCAmelCase   )


						@property
						def 					__lowercase	(							self			:     Dict   ):
										return ["input_ids", "attention_mask", "position_ids", "pixel_values"]


						@property
						def 					__lowercase	(							self			:     List[str]   ):
										warnings.warn(
										    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""  ,  lowerCAmelCase  ,  )
										return self.image_processor_class




						@property
						def 					__lowercase	(							self			:     Tuple   ):
										warnings.warn(
										    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""  ,  lowerCAmelCase  ,  )
										return self.image_processor





