


from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=1E-12      )->						List[str]:




      """simple docstring"""

      snake_case_ =   jnp.divide(emb_a.T      ,      jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE      ,      axis=1      )      ,      a_min=SCREAMING_SNAKE_CASE      )      ).T
      snake_case_ =   jnp.divide(emb_a.T      ,      jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE      ,      axis=1      )      ,      a_min=SCREAMING_SNAKE_CASE      )      ).T
      return jnp.matmul(SCREAMING_SNAKE_CASE      ,      norm_emb_a.T      )







class lowerCAmelCase_		(				nn.Module ):

 '''simple docstring'''

 __snake_case				=						4_2
 __snake_case				=						jnp.floataa



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   FlaxCLIPVisionModule(self.config.vision_config      )
       snake_case_ =   nn.Dense(self.config.projection_dim	,   use_bias=_UpperCAmelCase	,   dtype=self.dtype      )

       snake_case_ =   self.param('''concept_embeds'''	,   jax.nn.initializers.ones	,   (17, self.config.projection_dim)      )
       snake_case_ =   self.param(
           '''special_care_embeds'''	,   jax.nn.initializers.ones	,   (3, self.config.projection_dim)      )

       snake_case_ =   self.param('''concept_embeds_weights'''	,   jax.nn.initializers.ones	,   (17,)      )
       snake_case_ =   self.param('''special_care_embeds_weights'''	,   jax.nn.initializers.ones	,   (3,)      )





 def __call__(			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.vision_model(_UpperCAmelCase      )[1]
       snake_case_ =   self.visual_projection(_UpperCAmelCase      )

       snake_case_ =   jax_cosine_distance(_UpperCAmelCase	,   self.special_care_embeds      )
       snake_case_ =   jax_cosine_distance(_UpperCAmelCase	,   self.concept_embeds      )

       # increase this value to create a stronger `nfsw` filter
       # at the cost of increasing the possibility of filtering benign image inputs
       snake_case_ =   0.0

       snake_case_ =   special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
       snake_case_ =   jnp.round(_UpperCAmelCase	,   3      )
       snake_case_ =   jnp.any(special_scores > 0	,   axis=1	,   keepdims=_UpperCAmelCase      )
       # Use a lower threshold if an image has any special care concept
       snake_case_ =   is_special_care * 0.01

       snake_case_ =   cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
       snake_case_ =   jnp.round(_UpperCAmelCase	,   3      )
       snake_case_ =   jnp.any(concept_scores > 0	,   axis=1      )

       return has_nsfw_concepts







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						CLIPConfig
 __snake_case				=						"clip_input"
 __snake_case				=						FlaxStableDiffusionSafetyCheckerModule



 def __init__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase = None	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = jnp.floataa	,   _UpperCAmelCase = True	,   **_UpperCAmelCase	,   ):
       if input_shape is None:
             snake_case_ =   (1, 2_24, 2_24, 3)
       snake_case_ =   self.module_class(config=_UpperCAmelCase	,   dtype=_UpperCAmelCase	,   **_UpperCAmelCase      )
       super().__init__(_UpperCAmelCase	,   _UpperCAmelCase	,   input_shape=_UpperCAmelCase	,   seed=_UpperCAmelCase	,   dtype=_UpperCAmelCase	,   _do_init=_do_init      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase = None      ):
       # init input tensor
       snake_case_ =   jax.random.normal(_UpperCAmelCase	,   _UpperCAmelCase      )

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

       snake_case_ =   self.module.init(_UpperCAmelCase	,   _UpperCAmelCase      )['''params''']

       return random_params





 def __call__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase = None	,   ):
       snake_case_ =   jnp.transpose(_UpperCAmelCase	,   (0, 2, 3, 1)      )

       return self.module.apply(
           {'''params''': params or self.params}	,   jnp.array(_UpperCAmelCase	,   dtype=jnp.floataa      )	,   rngs={}	,   )




from typing import List, Optional, Union

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


UpperCAmelCase    =					logging.get_logger(__name__)

UpperCAmelCase    =					{
    """huggingface/informer-tourism-monthly""": (
        """https://huggingface.co/huggingface/informer-tourism-monthly/resolve/main/config.json"""
    ),
    # See all Informer models at https://huggingface.co/models?filter=informer
}







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						"informer"
 __snake_case				=						{
     "hidden_size": "d_model",
     "num_attention_heads": "encoder_attention_heads",
     "num_hidden_layers": "encoder_layers",
 }



 def __init__(			self	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = "student_t"	,   _UpperCAmelCase = "nll"	,   _UpperCAmelCase = 1	,   _UpperCAmelCase = None	,   _UpperCAmelCase = "mean"	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = 64	,   _UpperCAmelCase = 32	,   _UpperCAmelCase = 32	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = True	,   _UpperCAmelCase = "gelu"	,   _UpperCAmelCase = 0.05	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 1_00	,   _UpperCAmelCase = 0.02	,   _UpperCAmelCase=True	,   _UpperCAmelCase = "prob"	,   _UpperCAmelCase = 5	,   _UpperCAmelCase = True	,   **_UpperCAmelCase	,   ):
       # time series specific configuration
       snake_case_ =   prediction_length
       snake_case_ =   context_length or prediction_length
       snake_case_ =   distribution_output
       snake_case_ =   loss
       snake_case_ =   input_size
       snake_case_ =   num_time_features
       snake_case_ =   lags_sequence if lags_sequence is not None else [1, 2, 3, 4, 5, 6, 7]
       snake_case_ =   scaling
       snake_case_ =   num_dynamic_real_features
       snake_case_ =   num_static_real_features
       snake_case_ =   num_static_categorical_features

       # set cardinality
       if cardinality and num_static_categorical_features > 0:
             if len(_UpperCAmelCase      ) != num_static_categorical_features:
                   raise ValueError(
                       '''The cardinality should be a list of the same length as `num_static_categorical_features`'''      )
             snake_case_ =   cardinality
       else:
             snake_case_ =   [0]

       # set embedding_dimension
       if embedding_dimension and num_static_categorical_features > 0:
             if len(_UpperCAmelCase      ) != num_static_categorical_features:
                   raise ValueError(
                       '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''      )
             snake_case_ =   embedding_dimension
       else:
             snake_case_ =   [min(50	,   (cat + 1) // 2      ) for cat in self.cardinality]

       snake_case_ =   num_parallel_samples

       # Transformer architecture configuration
       snake_case_ =   input_size * len(self.lags_sequence      ) + self._number_of_features
       snake_case_ =   d_model
       snake_case_ =   encoder_attention_heads
       snake_case_ =   decoder_attention_heads
       snake_case_ =   encoder_ffn_dim
       snake_case_ =   decoder_ffn_dim
       snake_case_ =   encoder_layers
       snake_case_ =   decoder_layers

       snake_case_ =   dropout
       snake_case_ =   attention_dropout
       snake_case_ =   activation_dropout
       snake_case_ =   encoder_layerdrop
       snake_case_ =   decoder_layerdrop

       snake_case_ =   activation_function
       snake_case_ =   init_std

       snake_case_ =   use_cache

       # Informer
       snake_case_ =   attention_type
       snake_case_ =   sampling_factor
       snake_case_ =   distil

       super().__init__(is_encoder_decoder=_UpperCAmelCase	,   **_UpperCAmelCase      )





 @property
 def 		UpperCamelCase__  (			self      ):
       return (
           sum(self.embedding_dimension      )
           + self.num_dynamic_real_features
           + self.num_time_features
           + self.num_static_real_features
           + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
       )