



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 snake_case__  (			__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=1E-12 )		->							List[Any]:
		UpperCAmelCase_						=				jnp.divide(emb_a.T			,     jnp.clip(jnp.linalg.norm(__SCREAMING_SNAKE_CASE			,     axis=1 )			,     a_min=__SCREAMING_SNAKE_CASE ) ).T
		UpperCAmelCase_						=				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'''







					lowerCAmelCase_				:    CLIPConfig
					lowerCAmelCase_				:    jnp.dtype									=  jnp.floataa



					def    A__		(			self     ):
							UpperCAmelCase_						=				FlaxCLIPVisionModule(self.config.vision_config     )
							UpperCAmelCase_						=				nn.Dense(self.config.projection_dim    ,	use_bias=lowerCAmelCase    ,	dtype=self.dtype     )

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

							UpperCAmelCase_						=				self.param("concept_embeds_weights"    ,	jax.nn.initializers.ones    ,	(17,)     )
							UpperCAmelCase_						=				self.param("special_care_embeds_weights"    ,	jax.nn.initializers.ones    ,	(3,)     )



					def __call__(			self    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				self.vision_model(lowerCAmelCase     )[1]
							UpperCAmelCase_						=				self.visual_projection(lowerCAmelCase     )

							UpperCAmelCase_						=				jax_cosine_distance(lowerCAmelCase    ,	self.special_care_embeds     )
							UpperCAmelCase_						=				jax_cosine_distance(lowerCAmelCase    ,	self.concept_embeds     )

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

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

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

							return has_nsfw_concepts





class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					lowerCAmelCase_				:    Union[str, Any]									=  CLIPConfig
					lowerCAmelCase_				:    int									=  'clip_input'
					lowerCAmelCase_				:    Tuple									=  FlaxStableDiffusionSafetyCheckerModule



					def __init__(			self    ,	lowerCAmelCase    ,	lowerCAmelCase = None    ,	lowerCAmelCase = 0    ,	lowerCAmelCase = jnp.floataa    ,	lowerCAmelCase = True    ,	**lowerCAmelCase    ,	):
							if input_shape is None:
									UpperCAmelCase_						=				(1, 224, 224, 3)
							UpperCAmelCase_						=				self.module_class(config=lowerCAmelCase    ,	dtype=lowerCAmelCase    ,	**lowerCAmelCase     )
							super().__init__(lowerCAmelCase    ,	lowerCAmelCase    ,	input_shape=lowerCAmelCase    ,	seed=lowerCAmelCase    ,	dtype=lowerCAmelCase    ,	_do_init=_do_init     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase = None     ):
							# init input tensor
							UpperCAmelCase_						=				jax.random.normal(lowerCAmelCase    ,	lowerCAmelCase     )

							UpperCAmelCase_					,						UpperCAmelCase_						=				jax.random.split(lowerCAmelCase     )
							UpperCAmelCase_						=				{"params": params_rng, "dropout": dropout_rng}

							UpperCAmelCase_						=				self.module.init(lowerCAmelCase    ,	lowerCAmelCase     )["params"]

							return random_params



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

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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE     				=						{"configuration_fnet": ["FNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "FNetConfig"]}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						["FNetTokenizer"]

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						["FNetTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						[
			    "FNET_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "FNetForMaskedLM",
			    "FNetForMultipleChoice",
			    "FNetForNextSentencePrediction",
			    "FNetForPreTraining",
			    "FNetForQuestionAnswering",
			    "FNetForSequenceClassification",
			    "FNetForTokenClassification",
			    "FNetLayer",
			    "FNetModel",
			    "FNetPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_fnet import FNetTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_fnet_fast import FNetTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_fnet import (
						    FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
						    FNetForMaskedLM,
						    FNetForMultipleChoice,
						    FNetForNextSentencePrediction,
						    FNetForPreTraining,
						    FNetForQuestionAnswering,
						    FNetForSequenceClassification,
						    FNetForTokenClassification,
						    FNetLayer,
						    FNetModel,
						    FNetPreTrainedModel,
						)


else:
			import sys

			SCREAMING_SNAKE_CASE     				=						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

